from contextlib import asynccontextmanager from typing import List, Optional from uuid import uuid4 import uvicorn from fastapi import FastAPI, WebSocket, WebSocketDisconnect, HTTPException, Depends, status from fastapi.middleware.cors import CORSMiddleware from core.manager.websocket_manager import WebSocketManager from core.settings.settings import Settings from core.types.datamodels import ( ClientInfo, SubscriptionRequest, Subscription, NotificationRequest, Notification, BroadcastRequest, ) class WebSocketNotificationService: def __init__(self, settings: Optional[Settings] = None): self.settings = settings or Settings() self.manager = WebSocketManager(settings=self.settings) self.logger = self.manager.logger self.app = self._create_app() self._configure_middleware() self._register_routes() def _create_app(self) -> FastAPI: return FastAPI( title=self.settings.APP_NAME, description="Serivce de notification WebSocket et de real-time data", version=self.settings.API_VERSION, lifespan=self._lifespan, docs_url="/docs" if self.settings.APP_DEBUG else None, redoc_url="/redoc" if self.settings.APP_DEBUG else None, websocket_ping_interval=self.settings.WEBSOCKET_PING_INTERVAL, websocket_timeout=self.settings.WEBSOCKET_TIMEOUT, openapi_tags=[ { "name": "health", "description": "Santé et vérification du système" }, { "name": "websocket", "description": "Gestion de connexions WebSocket" }, { "name": "subscriptions", "description": "Gestion des subscriptions" }, { "name": "notifications", "description": "Envoie et gestion des notifications" }, { "name": "monitoring", "description": "Analytics et statistiques" } ] ) def _configure_middleware(self): self.app.add_middleware( CORSMiddleware, allow_origins=[self.settings.WS_CORS_ORIGINS], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) @asynccontextmanager async def _lifespan(self, _app: FastAPI): self.logger.info("🚀 Starting WebSocket notification service") try: await self.manager.initialize() yield finally: await self.manager.shutdown() self.logger.info("🛑 WebSocket notification service stopped") def _register_routes(self): self._register_health_routes() self._register_websocket_routes() self._register_subscription_routes() self._register_notification_routes() self._register_monitoring_routes() def _register_health_routes(self): @self.app.get("/", tags=["health"]) async def root(): return { "service": self.settings.APP_NAME, "version": self.settings.API_VERSION, "status": "healthy", "timestamp": self.manager.get_current_timestamp() } @self.app.get("/health", tags=["health"]) async def health_check(): return { "status": "healthy", "connections": len(self.manager.active_connections), "uptime": self.manager.get_uptime(), "memory_usage": self.manager.get_memory_usage() } def _register_websocket_routes(self): @self.app.websocket("/ws/{user_id}") async def websocket_endpoint(websocket: WebSocket, user_id: str): session_id = "" try: client_info = ClientInfo( session_id=str(uuid4()), user_id=user_id, user_roles=self._extract_roles(websocket), ip_address=self._get_client_ip(websocket), user_agent=self._get_user_agent(websocket) ) session_id = await self.manager.connect(websocket, client_info) await self._handle_websocket_messages(websocket, session_id) except WebSocketDisconnect: self.logger.info(f"WebSocket connection closed: {session_id}") except Exception as e: self.logger.error(f"WebSocket error: {e}") await self._handle_websocket_error(websocket, e) finally: self.manager.disconnect(session_id) def _register_subscription_routes(self): @self.app.post("/subscriptions/{session_id}", tags=["subscriptions"]) async def add_subscription( subscription_request: SubscriptionRequest, session_id: str = Depends(self._validate_session), ): try: subscription = Subscription(**subscription_request.model_dump()) if not self.manager.add_subscription(session_id, subscription): raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail="Failed to add subscription" ) return { "success": True, "message": "Subscription added successfully", "subscription_id": subscription.id, "timestamp": self.manager.get_current_timestamp() } except Exception as e: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail=str(e) ) @self.app.delete("/subscriptions/{session_id}/{subscription_id}", tags=["subscriptions"]) async def remove_subscription( subscription_id: str, session_id: str = Depends(self._validate_session), ): if not self.manager.remove_subscription(session_id, subscription_id): raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail="Subscription not found" ) return { "success": True, "message": "Subscription removed successfully", "timestamp": self.manager.get_current_timestamp() } @self.app.get("/subscriptions/{session_id}", tags=["subscriptions"]) async def get_subscriptions( session_id: str = Depends(self._validate_session), ): """Get all subscriptions for a session""" subscriptions = self.manager.get_subscriptions(session_id) return { "success": True, "subscriptions": subscriptions, "count": len(subscriptions), "timestamp": self.manager.get_current_timestamp() } def _register_notification_routes(self): @self.app.post("/notifications/user/{user_id}", tags=["notifications"]) async def send_user_notification( user_id: str, notification_request: NotificationRequest, ): try: notification = Notification(**notification_request.model_dump()) result = await self.manager.send_to_user(user_id, notification) return { "success": True, "message": f"Notification sent to user {user_id}", "notification_id": notification.id, "delivered_to": result.get("delivered_to", 0), "timestamp": self.manager.get_current_timestamp() } except Exception as e: raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=str(e) ) @self.app.post("/notifications/role/{role}", tags=["notifications"]) async def send_role_notification( role: str, notification_request: NotificationRequest, ): try: notification = Notification(**notification_request.model_dump()) result = await self.manager.send_to_role(role, notification) return { "success": True, "message": f"Notification sent to role {role}", "notification_id": notification.id, "delivered_to": result.get("delivered_to", 0), "timestamp": self.manager.get_current_timestamp() } except Exception as e: raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=str(e) ) @self.app.post("/notifications/broadcast", tags=["notifications"]) async def broadcast_notification( notification_request: NotificationRequest, ): try: notification = Notification(**notification_request.model_dump()) result = await self.manager.broadcast_notification(notification) return { "success": True, "message": "Notification broadcasted", "notification_id": notification.id, "delivered_to": result.get("delivered_to", 0), "timestamp": self.manager.get_current_timestamp() } except Exception as e: raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=str(e) ) @self.app.post("/notifications/send", tags=["notifications"]) async def send_targeted_notification( request: BroadcastRequest, ): try: notification = Notification(**request.notification.model_dump()) if request.target_type == "user": result = await self.manager.send_to_user(request.target_value, notification) elif request.target_type == "role": result = await self.manager.send_to_role(request.target_value, notification) elif request.target_type == "global": result = await self.manager.broadcast_notification(notification) else: raise ValueError("Invalid target type") return { "success": True, "message": f"Notification sent to {request.target_type}: {request.target_value}", "notification_id": notification.id, "delivered_to": result.get("delivered_to", 0), "timestamp": self.manager.get_current_timestamp() } except Exception as e: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail=str(e) ) def _register_monitoring_routes(self): @self.app.get("/notifications/history", tags=["monitoring"]) async def get_notification_history( limit: int = 100, offset: int = 0, ): history = self.manager.get_notification_history(limit, offset) return { "success": True, "history": history, "count": len(history), "limit": limit, "offset": offset, "timestamp": self.manager.get_current_timestamp() } @self.app.get("/stats", tags=["monitoring"]) async def get_service_stats(): stats = self.manager.get_stats() return { "success": True, "stats": stats, "timestamp": self.manager.get_current_timestamp() } @self.app.get("/clients", tags=["monitoring"]) async def get_connected_clients(): clients = [ { "session_id": session_id, "user_id": info.user_id, "user_roles": info.user_roles, "ip_address": info.ip_address, "connected_at": info.connected_at.isoformat(), "subscriptions_count": len(self.manager.subscriptions.get(session_id, [])), "last_activity": info.last_activity.isoformat() if hasattr(info, 'last_activity') else None } for session_id, info in self.manager.client_info.items() ] return { "success": True, "clients": clients, "total_connections": len(clients), "timestamp": self.manager.get_current_timestamp() } # Helper methods async def _validate_session(self, session_id: str) -> str: if session_id not in self.manager.active_connections: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail="Session not found or already closed" ) return session_id @staticmethod def _extract_roles(websocket: WebSocket) -> List[str]: roles_param = websocket.query_params.get("roles", "") return roles_param.split(",") if roles_param else [] @staticmethod def _get_client_ip(websocket: WebSocket) -> str: return websocket.client.host if websocket.client else "unknown" @staticmethod def _get_user_agent(websocket: WebSocket) -> str: return websocket.headers.get("user-agent", "unknown") async def _handle_websocket_messages(self, websocket: WebSocket, session_id: str): while True: try: message = await websocket.receive_text() await self._process_websocket_message(websocket, session_id, message) except WebSocketDisconnect: raise except Exception as e: self.logger.error(f"Error processing WebSocket message: {e}") await websocket.send_text(f"Error: {str(e)}") async def _process_websocket_message(self, websocket: WebSocket, session_id: str, message: str): if message == "ping": await websocket.send_text("pong") elif message.startswith("subscribe:"): topic = message.replace("subscribe:", "") subscription = Subscription(type="topic", target=topic) self.manager.add_subscription(session_id, subscription) await websocket.send_text(f"subscribed:{topic}") elif message.startswith("unsubscribe:"): topic = message.replace("unsubscribe:", "") await websocket.send_text(f"unsubscribed:{topic}") else: await self._handle_custom_message(websocket, session_id, message) async def _handle_custom_message(self, websocket: WebSocket, session_id: str, message: str): """Handle custom WebSocket messages""" pass @staticmethod async def _handle_websocket_error(websocket: WebSocket, error: Exception): """Handle WebSocket errors gracefully""" try: await websocket.send_text(f"Error: {str(error)}") except: raise def run(self): uvicorn.run( self.app, host=self.settings.APP_HOST, port=self.settings.APP_PORT, reload=self.settings.APP_DEBUG, log_level="info", server_header=False, ) def create_app(settings: Optional[Settings] = None) -> FastAPI: service = WebSocketNotificationService(settings) return service.app def main(): service = WebSocketNotificationService() service.run() app = create_app() if __name__ == "__main__": main()