main.py

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)