Update app.py

app.py

@@ -1,23 +1,23 @@
 from typing import List, Dict, Optional
 from datetime import datetime, timedelta
 from fastapi import FastAPI, HTTPException, Query, Body, Request
-from pydantic import BaseModel, validator
+from pydantic import BaseModel, validator, root_validator
 import json
 import os
+from user_agents import parse
 
 app = FastAPI()
 
 # Data storage (in-memory)
 user_data: Dict[str, dict] = {}  # Key: IP address, Value: User entry
 
-# Data storage file
-DATA_FILE = "user_data.json"
-
 # --- Data Models ---
 class UserEntry(BaseModel):
     ip_address: str
-    device_type: str
-    timestamp: datetime
+    device_type: str = "N/A"
+    timestamp: datetime = datetime.now()
+    browser: str = "N/A"
+    OS: str = "N/A"
 
     @validator("ip_address")
     def validate_ip_address(cls, value):
@@ -33,36 +33,19 @@ class UserEntry(BaseModel):
                 raise ValueError("Invalid IP address value")
         return value
 
-# --- Helper Functions ---
-def load_data():
-    """Loads data from the JSON file into the in-memory store."""
-    global user_data
-    if os.path.exists(DATA_FILE):
-        try:
-            with open(DATA_FILE, "r") as f:
-                data = json.load(f)
-                # Convert timestamp strings back to datetime objects
-                user_data = {
-                    ip: {**entry, "timestamp": datetime.fromisoformat(entry["timestamp"])}
-                    for ip, entry in data.items()
-                }
-        except Exception as e:
-            print(f"Error loading data from {DATA_FILE}: {e}")
-            # Handle the error gracefully, maybe start with an empty dataset
-            user_data = {}
-
-def save_data():
-    """Saves the in-memory data to the JSON file."""
-    try:
-        with open(DATA_FILE, "w") as f:
-            # Convert datetime objects to ISO format for JSON serialization
-            serializable_data = {
-                ip: {**entry, "timestamp": entry["timestamp"].isoformat()}
-                for ip, entry in user_data.items()
-            }
-            json.dump(serializable_data, f, indent=4)
-    except Exception as e:
-        print(f"Error saving data to {DATA_FILE}: {e}")
+    @root_validator(pre=True)
+    def set_default_values(cls, values):
+        """Set default values for missing fields."""
+        defaults = {
+            "device_type": "N/A",
+            "browser": "N/A",
+            "OS": "N/A",
+            "timestamp": datetime.now()
+        }
+        for key, default_value in defaults.items():
+            if key not in values or values[key] is None:
+                values[key] = default_value
+        return values
 
 def clean_old_data():
     """Deletes data older than 7 days."""
@@ -71,18 +54,11 @@ def clean_old_data():
     ips_to_delete = [ip for ip, entry in user_data.items() if entry["timestamp"] < cutoff_time]
     for ip in ips_to_delete:
         del user_data[ip]
-    if ips_to_delete:
-        save_data()  # Save changes after deleting old data
-
-# Load data on startup
-load_data()
 
 # --- API Endpoints ---
 @app.post("/auto_entry/", response_model=UserEntry, status_code=201)
 async def create_auto_user_entry(
-    request: Request,
-    device_type: str = Body(...),
-    timestamp: Optional[datetime] = Body(None)
+    request: Request
 ):
     """
     Endpoint to automatically record user entry by extracting the IP address
@@ -94,20 +70,26 @@ async def create_auto_user_entry(
         if "x-forwarded-for" in request.headers:
             ip_address = request.headers["x-forwarded-for"].split(",")[0]
 
-        # If timestamp is not provided, use the current time
-        if not timestamp:
-            timestamp = datetime.now()
+        user_agent = request.headers.get("User-Agent", "N/A")
+        user_agent_parsed = parse(user_agent)
+    
+        device_type = "Mobile" if user_agent_parsed.is_mobile else "Tablet" if user_agent_parsed.is_tablet else "Desktop"
+        browser_name = user_agent_parsed.browser.family if user_agent_parsed else "N/A"
+        os_name = user_agent_parsed.os.family if user_agent_parsed else "N/A"
+
+        timestamp = datetime.now()
 
         # Create a UserEntry object
         entry_data = UserEntry(
             ip_address=ip_address,
             device_type=device_type,
-            timestamp=timestamp
+            timestamp=timestamp,
+            browser=browser_name,
+            OS=os_name
         )
 
         # Save the entry
         user_data[ip_address] = entry_data.dict()
-        save_data()
 
         return entry_data
 
@@ -116,19 +98,6 @@ async def create_auto_user_entry(
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Internal server error: {e}")
 
-@app.post("/entry/", response_model=UserEntry, status_code=201)
-async def create_user_entry(entry_data: UserEntry):
-    """Endpoint to record user entry."""
-    try:
-        entry_data.timestamp = datetime.now()
-        user_data[entry_data.ip_address] = entry_data.dict()
-        save_data()
-        return entry_data
-    except ValueError as ve:
-        raise HTTPException(status_code=400, detail=str(ve))
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Internal server error: {e}")
-
 @app.get("/analytics/")
 async def get_user_analytics(period: str = Query(..., enum=["last_hour", "last_day", "last_7_day"])):
     """Endpoint to get user analytics."""
@@ -178,19 +147,92 @@ async def import_user_data(data: Dict[str, dict] = Body(...)):
             try:
                 # Validate the imported entry
                 entry = UserEntry(**entry_dict)
-                entry.timestamp = datetime.fromisoformat(entry_dict["timestamp"]) # Ensure timestamp is datetime
+                entry.timestamp = datetime.fromisoformat(entry_dict.get("timestamp", datetime.now().isoformat())) # Ensure timestamp is datetime
                 user_data[ip] = entry.dict()
                 imported_count += 1
             except Exception as e:
                 print(f"Error importing entry for IP {ip}: {e}")  # Log individual import errors
-
-        save_data()
         return {"message": f"Successfully imported {imported_count} user entries."}
     except json.JSONDecodeError:
         raise HTTPException(status_code=400, detail="Invalid JSON format")
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Error importing data: {e}")
 
+# Data storage (in-memory)
+user_data: Dict[str, dict] = {}  # Key: IP address, Value: User entry
+poll_data: Dict[str, dict] = {}  # Key: Poll name, Value: Poll details
+poll_responses: Dict[str, Dict[str, int]] = {}  # Key: Poll name, Value: {IP: response}
+
+# --- Data Models ---
+class PollCreate(BaseModel):
+    poll_name: str
+    question: str
+    options: List[str]
+
+class PollEntry(BaseModel):
+    poll_name: str
+    response: int
+
+# --- API Endpoints ---
+@app.post("/poll/create/", status_code=201)
+async def create_poll(poll: PollCreate):
+    """Endpoint to create a new poll."""
+    if poll.poll_name in poll_data:
+        raise HTTPException(status_code=400, detail="Poll with this name already exists.")
+    poll_data[poll.poll_name] = {
+        "question": poll.question,
+        "options": poll.options,
+        "created_at": datetime.now()
+    }
+    poll_responses[poll.poll_name] = {}
+    return {"message": "Poll created successfully.", "poll_name": poll.poll_name}
+
+@app.post("/poll/entry/", status_code=201)
+async def create_poll_entry(request: Request, poll_entry: PollEntry):
+    """Endpoint to record a user's response to a poll."""
+    ip_address = request.client.host
+    if "x-forwarded-for" in request.headers:
+        ip_address = request.headers["x-forwarded-for"].split(",")[0]
+
+    if poll_entry.poll_name not in poll_data:
+        raise HTTPException(status_code=404, detail="Poll not found.")
+
+    if poll_entry.response < 1 or poll_entry.response > len(poll_data[poll_entry.poll_name]["options"]):
+        raise HTTPException(status_code=400, detail="Invalid response option.")
+
+    poll_responses[poll_entry.poll_name][ip_address] = poll_entry.response
+    return {"message": "Poll entry recorded successfully."}
+
+@app.get("/poll/analytics/")
+async def get_poll_analytics(poll_name: str = Query(..., description="Name of the poll")):
+    """Endpoint to get analytics for a specific poll."""
+    if poll_name not in poll_data:
+        raise HTTPException(status_code=404, detail="Poll not found.")
+
+    responses = poll_responses[poll_name]
+    response_counts = {option: 0 for option in range(1, len(poll_data[poll_name]["options"]) + 1)}
+    for response in responses.values():
+        response_counts[response] += 1
+
+    return {
+        "poll_name": poll_name,
+        "question": poll_data[poll_name]["question"],
+        "options": poll_data[poll_name]["options"],
+        "response_counts": response_counts,
+        "total_responses": len(responses)
+    }
+
+# --- Background Task for Poll Deletion ---
+async def scheduled_poll_cleanup():
+    """Periodically clean up old polls."""
+    while True:
+        now = datetime.now()
+        polls_to_delete = [poll_name for poll_name, poll in poll_data.items() if (now - poll["created_at"]).days >= 7]
+        for poll_name in polls_to_delete:
+            del poll_data[poll_name]
+            del poll_responses[poll_name]
+        await asyncio.sleep(60 * 60 * 24)  # Clean every 24 hours
+    
 # --- Background Task (Optional, for regular cleanup) ---
 async def scheduled_cleanup():
     """Periodically clean up old data."""
@@ -206,6 +248,7 @@ from fastapi import BackgroundTasks
 async def startup_event():
     # You can uncomment this to run the background task
     asyncio.create_task(scheduled_cleanup())
+    asyncio.create_task(scheduled_poll_cleanup())
     pass
 
 # --- Error Handling (Advanced - using exception handlers) ---