import uvicorn
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from config import setup_app, agent, logger, patients_collection, analysis_collection, users_collection, chats_collection, notifications_collection
from endpoints import create_router
from fastapi import WebSocket, WebSocketDisconnect

# Create the FastAPI app
app = FastAPI(title="TxAgent API", version="2.6.0")

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

@app.websocket("/ws/notifications")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    try:
        while True:
            # Keep connection alive
            await websocket.receive_text()
    except WebSocketDisconnect:
        logger.info("Client disconnected")

# Setup the app (e.g., initialize globals, startup event)
setup_app(app)

# Create and include the router with dependencies
router = create_router(agent, logger, patients_collection, analysis_collection, users_collection, chats_collection, notifications_collection)
app.include_router(router)

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)