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)