Attempting to add websocket support

main.py

@@ -14,6 +14,7 @@ import schemas
 from starlette.responses import Response
 from config import settings # Import settings to get the secret key
 from listener_manager import ListenerManager
+from websocket_manager import WebSocketManager
 
 # --- Absolute Path Configuration ---
 # Get the absolute path of the directory where this file is located
@@ -25,6 +26,7 @@ TEMPLATES_DIR = os.path.join(BASE_DIR, "templates")
 async def lifespan(app: FastAPI):
     # This code runs on startup
     print("Application startup: Creating database tables...")
+    app.state.websocket_manager = WebSocketManager()
     app.state.listener_manager = ListenerManager()
     models.Base.metadata.create_all(bind=engine)
     print("Application startup: Database tables created.")
@@ -34,7 +36,7 @@ async def lifespan(app: FastAPI):
     users = db.query(models.User).all()
     db.close()
     for user in users:
-        await app.state.listener_manager.start_listener_for_user(user)
+        await app.state.listener_manager.start_listener_for_user(user, app.state.websocket_manager)
 
     yield
     
@@ -170,4 +172,16 @@ async def get_custom_css(theme_id: int, db: Session = Depends(auth.get_db)):
     if not theme:
         raise HTTPException(status_code=404, detail="Custom theme not found")
     
-    return Response(content=theme.css_content, media_type="text/css")
+    return Response(content=theme.css_content, media_type="text/css")
+
+@app.websocket("/ws/{user_id}")
+async def websocket_endpoint(websocket: WebSocket, user_id: int):
+    manager = websocket.app.state.websocket_manager
+    await manager.connect(user_id, websocket)
+    try:
+        while True:
+            # Keep the connection alive
+            await websocket.receive_text()
+    except Exception:
+        manager.disconnect(user_id, websocket)
+        print(f"WebSocket for user {user_id} disconnected.")