app/main.py

import uvicorn
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
from contextlib import asynccontextmanager

from app.config import env
from app.database import connect_db, close_db
from app.routers import assignments, papers
from app.services.ws_service import manager

@asynccontextmanager
async def lifespan(app: FastAPI):
    # Startup
    await connect_db()
    yield
    # Shutdown
    await close_db()

app = FastAPI(lifespan=lifespan, title="VedaAI Backend")

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

# Health check
@app.get("/health")
async def health():
    from datetime import datetime
    return {"status": "ok", "timestamp": datetime.utcnow().isoformat()}

# Routers
app.include_router(assignments.router, prefix="/api/assignments", tags=["Assignments"])
app.include_router(papers.router, prefix="/api/papers", tags=["Papers"])

# WebSocket Endpoint
@app.websocket("/ws/{room_id}")
async def websocket_endpoint(websocket: WebSocket, room_id: str):
    await manager.connect(websocket, room_id)
    try:
        while True:
            # Just keep the connection open, maybe accept ping messages
            data = await websocket.receive_text()
    except WebSocketDisconnect:
        manager.disconnect(websocket, room_id)

if __name__ == "__main__":
    print(f"VedaAI backend running on port {env.PORT}")
    print(f"   AI Provider: groq")
    print(f"   CORS origins: {env.CORS_ORIGINS}")
    print(f"   Health check: http://localhost:{env.PORT}/health")
    uvicorn.run("app.main:app", host="0.0.0.0", port=env.PORT, reload=True)