"""WebSocket connection manager for real-time updates.""" from fastapi import WebSocket from typing import List, Dict, Any import json import asyncio import logging logger = logging.getLogger(__name__) class ConnectionManager: """Manages WebSocket connections for broadcasting events.""" def __init__(self): self.active_connections: List[WebSocket] = [] async def connect(self, websocket: WebSocket): """Accept a new WebSocket connection.""" await websocket.accept() self.active_connections.append(websocket) logger.info(f"New WebSocket connection. Total connections: {len(self.active_connections)}") def disconnect(self, websocket: WebSocket): """Remove a WebSocket connection.""" if websocket in self.active_connections: self.active_connections.remove(websocket) logger.info(f"WebSocket disconnected. Total connections: {len(self.active_connections)}") async def send_personal_message(self, message: Dict[str, Any], websocket: WebSocket): """Send a message to a specific connection.""" try: await websocket.send_text(json.dumps(message)) except Exception as e: logger.error(f"Error sending personal message: {e}") async def broadcast(self, message: Dict[str, Any]): """Broadcast a message to all connected clients.""" if not self.active_connections: return # Convert message to JSON string message_str = json.dumps(message) # Send to all connections disconnected = [] for connection in self.active_connections: try: await connection.send_text(message_str) except Exception as e: logger.error(f"Error broadcasting to connection: {e}") disconnected.append(connection) # Clean up disconnected connections for connection in disconnected: self.disconnect(connection) async def broadcast_event(self, event_data: Dict[str, Any], statistics_data: Dict[str, Any]): """Broadcast a new event with updated statistics.""" message = { "type": "event_created", "event": event_data, "statistics": statistics_data, } await self.broadcast(message) async def broadcast_stats_reset(self): """Broadcast a stats reset notification.""" message = { "type": "stats_reset", "message": "Statistics have been reset", } await self.broadcast(message) async def broadcast_stats_update(self, statistics_data: Dict[str, Any]): """Broadcast updated statistics.""" message = { "type": "stats_updated", "statistics": statistics_data, } await self.broadcast(message) def get_connection_count(self) -> int: """Get the number of active connections.""" return len(self.active_connections) # Global connection manager instance manager = ConnectionManager()