import asyncio from collections import deque from datetime import datetime, timezone from typing import Dict, List, Any, Deque, Optional import psutil from fastapi import WebSocket from core.settings.settings import Settings from core.types.datamodels import ClientInfo, Subscription, Notification from core.utils.utils import current_utc_iso class WebSocketManager: def __init__(self, settings: Settings) -> None: self.settings = settings self.logger = self.settings.logger self.active_connections: Dict[str, WebSocket] = {} self.client_info: Dict[str, ClientInfo] = {} self.subscriptions: Dict[str, List[Subscription]] = {} self.notification_history: Deque[Notification] = deque( maxlen=self.settings.get("MAX_NOTIFICATION_HISTORY", 1000) ) self._start_time: Optional[datetime] = None self._is_initialized: bool = False self._shutdown_event: Optional[asyncio.Event] = None self._background_tasks: List[asyncio.Task] = [] self._total_messages_sent: int = 0 self._total_connections_handled: int = 0 self._connection_errors: int = 0 async def initialize(self) -> None: if self._is_initialized: self.logger.warning("WebSocketManager already initialized") return self._start_time = datetime.now(timezone.utc) self._shutdown_event = asyncio.Event() self._is_initialized = True self._background_tasks = [ asyncio.create_task(self._cleanup_expired_notifications()), asyncio.create_task(self._heartbeat_checker()), asyncio.create_task(self._stats_collector()) ] self.logger.info("WebSocketManager initialized successfully") async def shutdown(self) -> None: if not self._is_initialized: return self.logger.info("Shutting down WebSocketManager...") if self._shutdown_event: self._shutdown_event.set() for task in self._background_tasks: if not task.done(): task.cancel() if self._background_tasks: await asyncio.gather(*self._background_tasks, return_exceptions=True) await self._close_all_connections() self.active_connections.clear() self.client_info.clear() self.subscriptions.clear() self.notification_history.clear() self._is_initialized = False self.logger.info("WebSocketManager shutdown completed") async def connect(self, websocket: WebSocket, client: ClientInfo) -> str: await websocket.accept() session_id = client.session_id self.active_connections[session_id] = websocket self.client_info[session_id] = client self.subscriptions[session_id] = [] self._total_connections_handled += 1 self.logger.info(f"Nouvelle connexion: {session_id}") await self._send( websocket, { "session_id": session_id, "type": "connection_established", "title": "Connexion réussie", "message": "Connexion au service de notifications établie", "timestamp": current_utc_iso(), }, ) return session_id def disconnect(self, session_id: str) -> None: if session_id in self.active_connections: del self.active_connections[session_id] del self.client_info[session_id] del self.subscriptions[session_id] self.logger.info(f"Connexion fermée: {session_id}") async def send_personal_message( self, session_id: str, payload: Dict[str, Any] ) -> None: if ws := self.active_connections.get(session_id): try: await self._send(ws, payload) self._total_messages_sent += 1 except RuntimeError as e: self.logger.error(f"Erreur d'envoi à {session_id}: {e}") self._connection_errors += 1 self.disconnect(session_id) async def broadcast_notification(self, notification: Notification) -> Dict[str, Any]: if notification.is_expired: return {"delivered_to": 0} self._add_to_history(notification) payload = self._format_notification(notification) delivered_count = 0 for session_id, subscriptions in self.subscriptions.items(): if not (client := self.client_info.get(session_id)): continue if self._should_receive(subscriptions, notification, client): await self.send_personal_message(session_id, payload) delivered_count += 1 return {"delivered_to": delivered_count} async def send_to_user(self, user_id: str, notification: Notification) -> Dict[str, Any]: if notification.is_expired: return {"delivered_to": 0} self._add_to_history(notification) payload = self._format_notification(notification) delivered_count = 0 for session_id, client in self.client_info.items(): if client.user_id == user_id: await self.send_personal_message(session_id, payload) delivered_count += 1 return {"delivered_to": delivered_count} async def send_to_role(self, role: str, notification: Notification) -> Dict[str, Any]: if notification.is_expired: return {"delivered_to": 0} self._add_to_history(notification) payload = self._format_notification(notification) delivered_count = 0 for session_id, client in self.client_info.items(): if role in client.user_roles: await self.send_personal_message(session_id, payload) delivered_count += 1 return {"delivered_to": delivered_count} def add_subscription(self, session_id: str, subscription: Subscription) -> bool: if not (sub_list := self.subscriptions.get(session_id)): return False for existing in sub_list: if existing.id == subscription.id: return False sub_list.append(subscription) return True def update_subscription(self, session_id: str, subscription: Subscription) -> bool: if not (sub_list := self.subscriptions.get(session_id)): return False for i, existing in enumerate(sub_list): if existing.id == subscription.id: sub_list[i] = subscription return True return False def remove_subscription(self, session_id: str, subscription_id: str) -> bool: if not (sub_list := self.subscriptions.get(session_id)): return False initial_count = len(sub_list) self.subscriptions[session_id] = [ sub for sub in sub_list if sub.id != subscription_id ] return len(self.subscriptions[session_id]) < initial_count def get_subscriptions(self, session_id: str) -> List[Subscription]: return self.subscriptions.get(session_id, [])[:] def get_notification_history(self, limit: int = 100, offset: int = 0) -> List[Notification]: history = [n for n in list(self.notification_history) if not n.is_expired] return history[offset:offset + limit] def get_stats(self) -> Dict[str, Any]: memory_info = self.get_memory_usage() return { "active_connections": len(self.active_connections), "total_subscriptions": sum(len(s) for s in self.subscriptions.values()), "total_notifications": len(self.notification_history), "total_messages_sent": self._total_messages_sent, "total_connections_handled": self._total_connections_handled, "connection_errors": self._connection_errors, "uptime_seconds": self.get_uptime(), "memory_usage": memory_info, "timestamp": current_utc_iso(), } @staticmethod def get_current_timestamp() -> str: return current_utc_iso() def get_uptime(self) -> str: if not self._start_time: return self.format_uptime(0) return self.format_uptime((datetime.now(timezone.utc) - self._start_time).total_seconds()) def get_memory_usage(self) -> Dict[str, Any]: try: process = psutil.Process() memory_info = process.memory_info() return { "resident_set_size_mb": self._format_size(memory_info.rss), "virtual_memory_size_mb": self._format_size(memory_info.vms), "percent": self._format_size(process.memory_percent()), "available_mb": self._format_size(psutil.virtual_memory().available), "total_mb": self._format_size(psutil.virtual_memory().total), } except Exception as e: self.logger.error(f"Error getting memory usage: {e}") return { "resident_set_size_mb": 0, "virtual_memory_size_mb": 0, "percent": 0, "available_mb": 0, "total_mb": 0, "error": str(e) } async def _cleanup_expired_notifications(self) -> None: while not self._shutdown_event.is_set(): try: active_notifications = [ n for n in self.notification_history if not n.is_expired ] self.notification_history.clear() self.notification_history.extend(active_notifications) await asyncio.sleep(self.settings.get("CLEANUP_INTERVAL", 300)) # 5 minutes except asyncio.CancelledError: break except Exception as e: self.logger.error(f"Error in cleanup task: {e}") await asyncio.sleep(60) async def _heartbeat_checker(self) -> None: while not self._shutdown_event.is_set(): try: disconnected_sessions = [] for session_id, websocket in self.active_connections.items(): try: await websocket.send_json( {"session_id": session_id, "type": "ping", "timestamp": current_utc_iso()}) except Exception as e: self.logger.warning(f"Connection {session_id} appears to be dead: {e}") disconnected_sessions.append(session_id) for session_id in disconnected_sessions: self.disconnect(session_id) await asyncio.sleep(self.settings.get("HEARTBEAT_INTERVAL", 30)) # 30 seconds except asyncio.CancelledError: break except Exception as e: self.logger.error(f"Error in heartbeat checker: {e}") await asyncio.sleep(60) async def _stats_collector(self) -> None: while not self._shutdown_event.is_set(): try: _ = self.get_stats() await asyncio.sleep(self.settings.get("STATS_INTERVAL", 60)) # 1 minute except asyncio.CancelledError: break except Exception as e: self.logger.error(f"Error in stats collector: {e}") await asyncio.sleep(60) async def _close_all_connections(self) -> None: if not self.active_connections: return self.logger.info(f"Closing {len(self.active_connections)} active connections...") shutdown_message = { "type": "service_shutdown", "title": "Service en cours d'arrêt", "message": "Le service de notifications va être arrêté", "timestamp": current_utc_iso() } close_tasks = [] for session_id, websocket in self.active_connections.items(): try: await self._send(websocket, shutdown_message) close_tasks.append(websocket.close()) except Exception as e: self.logger.error(f"Error closing connection {session_id}: {e}") if close_tasks: await asyncio.gather(*close_tasks, return_exceptions=True) async def _send(self, websocket: WebSocket, payload: Dict[str, Any]) -> None: try: await websocket.send_json(payload) except Exception as e: self.logger.error(f"Failed to send message: {e}") raise RuntimeError(f"WebSocket send failed: {e}") def _add_to_history(self, notification: Notification) -> None: if not notification.is_expired: self.notification_history.append(notification) @staticmethod def _should_receive( subscriptions: List[Subscription], notification: Notification, client: ClientInfo, ) -> bool: return any( sub.matches(notification, client.user_id, client.user_roles) for sub in subscriptions ) @staticmethod def _format_notification(notification: Notification) -> Dict[str, Any]: return {"type": "notification", **notification.model_dump(exclude_unset=True)} @staticmethod def _format_size(bytes_value) -> str: if bytes_value < 1024: return f"{bytes_value} B" elif bytes_value < 1024 ** 2: return f"{bytes_value / 1024:.2f} KB" elif bytes_value < 1024 ** 3: return f"{bytes_value / (1024 ** 2):.2f} MB" else: return f"{bytes_value / (1024 ** 3):.2f} GB" @staticmethod def format_uptime(seconds: float) -> str: minutes, sec = divmod(int(seconds), 60) hours, minutes = divmod(minutes, 60) days, hours = divmod(hours, 24) parts = [] if days: parts.append(f"{days}d") if hours: parts.append(f"{hours}h") if minutes: parts.append(f"{minutes}m") if sec or not parts: parts.append(f"{sec}s") return " ".join(parts)