src/main.py

from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Depends, HTTPException, status, Request
from fastapi.staticfiles import StaticFiles
from fastapi.responses import HTMLResponse, JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from sqlalchemy.orm import Session
import json
import asyncio
from typing import Optional

from database import engine, get_db, Base
from models import User, Message, Conversation, ConversationMember, GroupMessage, Group, GroupMember
from auth import router as auth_router, get_current_user_ws, get_current_user
from chat import router as chat_router
from ws_manager import manager

Base.metadata.create_all(bind=engine)

app = FastAPI(title="SecureChat")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

app.include_router(auth_router, prefix="/api/auth")
app.include_router(chat_router, prefix="/api/chat")

app.mount("/static", StaticFiles(directory="static"), name="static")

@app.get("/", response_class=HTMLResponse)
async def root():
    with open("static/index.html") as f:
        return f.read()

@app.websocket("/ws/{token}")
async def websocket_endpoint(websocket: WebSocket, token: str, db: Session = Depends(get_db)):
    user = await get_current_user_ws(token, db)
    if not user:
        await websocket.close(code=4001)
        return

    await manager.connect(websocket, user.id)
    try:
        while True:
            data = await websocket.receive_text()
            msg = json.loads(data)
            msg_type = msg.get("type")

            if msg_type == "dm":
                await manager.handle_dm(msg, user, db)
            elif msg_type == "group":
                await manager.handle_group(msg, user, db)
            elif msg_type == "typing":
                await manager.handle_typing(msg, user)
            elif msg_type == "read":
                await manager.handle_read(msg, user, db)
            elif msg_type == "call_offer":
                await manager.handle_call_offer(msg, user)
            elif msg_type == "call_answer":
                await manager.handle_call_answer(msg, user)
            elif msg_type == "ice_candidate":
                await manager.handle_ice_candidate(msg, user)
            elif msg_type == "call_end":
                await manager.handle_call_end(msg, user)
            elif msg_type == "video_toggle":
                await manager.handle_video_toggle(msg, user)

    except WebSocketDisconnect:
        manager.disconnect(user.id)
    except Exception as e:
        manager.disconnect(user.id)