Create main.py

main.py

@@ -0,0 +1,173 @@
+import contextlib
+import sqlite3
+import uuid
+from datetime import datetime
+from typing import List, Dict
+
+from fastapi import FastAPI, WebSocket, WebSocketDisconnect, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+
+# --- Configuration ---
+DB_FILE = "messaging_app.db"
+
+# --- Database Setup ---
+def init_db():
+    """Initialize the SQLite database safely."""
+    try:
+        conn = sqlite3.connect(DB_FILE)
+        cursor = conn.cursor()
+        cursor.execute("""
+            CREATE TABLE IF NOT EXISTS messages (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                room_id TEXT NOT NULL,
+                username TEXT NOT NULL,
+                content TEXT NOT NULL,
+                timestamp DATETIME DEFAULT CURRENT_TIMESTAMP
+            )
+        """)
+        cursor.execute("CREATE INDEX IF NOT EXISTS idx_room_id ON messages(room_id)")
+        conn.commit()
+        conn.close()
+        print("Database initialized successfully.")
+    except Exception as e:
+        print(f"Error initializing database: {e}")
+
+# --- Lifespan Management ---
+# This ensures DB init happens only when the app actually starts, 
+# preventing the "Import Error" you saw.
+@contextlib.asynccontextmanager
+async def lifespan(app: FastAPI):
+    # Startup
+    init_db()
+    yield
+    # Shutdown (Cleanup code goes here if needed)
+
+# --- Pydantic Models ---
+class Message(BaseModel):
+    room_id: str
+    username: str
+    content: str
+
+class RoomResponse(BaseModel):
+    room_id: str
+    link: str
+
+# --- Connection Manager ---
+class ConnectionManager:
+    def __init__(self):
+        self.active_connections: Dict[str, List[WebSocket]] = {}
+
+    async def connect(self, websocket: WebSocket, room_id: str):
+        await websocket.accept()
+        if room_id not in self.active_connections:
+            self.active_connections[room_id] = []
+        self.active_connections[room_id].append(websocket)
+
+    def disconnect(self, websocket: WebSocket, room_id: str):
+        if room_id in self.active_connections:
+            self.active_connections[room_id].remove(websocket)
+            if not self.active_connections[room_id]:
+                del self.active_connections[room_id]
+
+    async def broadcast(self, message: dict, room_id: str):
+        if room_id in self.active_connections:
+            for connection in self.active_connections[room_id]:
+                try:
+                    await connection.send_json(message)
+                except Exception:
+                    pass
+
+manager = ConnectionManager()
+
+# --- FastAPI App Initialization ---
+# Pass the lifespan handler here
+app = FastAPI(title="Room-Based Messaging API", lifespan=lifespan)
+
+# --- CORS Middleware ---
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# --- API Endpoints ---
+
+@app.get("/")
+async def root():
+    return {"status": "online", "message": "Messaging Backend is running"}
+
+@app.post("/api/room/create", response_model=RoomResponse)
+async def create_room():
+    room_id = str(uuid.uuid4())
+    return {"room_id": room_id, "link": f"/room/{room_id}"}
+
+@app.get("/api/room/{room_id}/history")
+async def get_room_history(room_id: str):
+    conn = sqlite3.connect(DB_FILE)
+    conn.row_factory = sqlite3.Row
+    cursor = conn.cursor()
+    try:
+        cursor.execute(
+            "SELECT username, content, timestamp FROM messages WHERE room_id = ? ORDER BY timestamp ASC",
+            (room_id,)
+        )
+        rows = cursor.fetchall()
+        history = [
+            {
+                "username": row["username"],
+                "content": row["content"],
+                "timestamp": row["timestamp"],
+                "type": "history"
+            }
+            for row in rows
+        ]
+        return {"room_id": room_id, "messages": history}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+    finally:
+        conn.close()
+
+# --- WebSocket Endpoint ---
+
+@app.websocket("/ws/{room_id}/{username}")
+async def websocket_endpoint(websocket: WebSocket, room_id: str, username: str):
+    await manager.connect(websocket, room_id)
+    
+    async def process_message(content: str):
+        message_payload = {
+            "username": username,
+            "content": content,
+            "timestamp": datetime.now().isoformat(),
+            "type": "live"
+        }
+        try:
+            conn = sqlite3.connect(DB_FILE)
+            cursor = conn.cursor()
+            cursor.execute(
+                "INSERT INTO messages (room_id, username, content) VALUES (?, ?, ?)",
+                (room_id, username, content)
+            )
+            conn.commit()
+            conn.close()
+        except Exception as e:
+            print(f"Database error: {e}")
+
+        await manager.broadcast(message_payload, room_id)
+
+    try:
+        while True:
+            data = await websocket.receive_text()
+            await process_message(data)
+    except WebSocketDisconnect:
+        manager.disconnect(websocket, room_id)
+        await manager.broadcast({
+            "username": "System",
+            "content": f"{username} left the chat",
+            "type": "system"
+        }, room_id)
+    except Exception as e:
+        print(f"Error: {e}")
+        manager.disconnect(websocket, room_id)