""" WebSocket streaming API endpoints """ import json import logging from typing import Dict, List, Optional, Any from datetime import datetime from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Depends, HTTPException, Query from fastapi.responses import JSONResponse from pydantic import BaseModel, Field from src.api.dependencies import ( get_stream_service, get_pose_service, get_current_user_ws, require_auth ) from src.api.websocket.connection_manager import ConnectionManager from src.services.stream_service import StreamService from src.services.pose_service import PoseService logger = logging.getLogger(__name__) router = APIRouter() # Initialize connection manager connection_manager = ConnectionManager() # Request/Response models class StreamSubscriptionRequest(BaseModel): """Request model for stream subscription.""" zone_ids: Optional[List[str]] = Field( default=None, description="Zones to subscribe to (all zones if not specified)" ) stream_types: List[str] = Field( default=["pose_data"], description="Types of data to stream" ) min_confidence: float = Field( default=0.5, ge=0.0, le=1.0, description="Minimum confidence threshold for streaming" ) max_fps: int = Field( default=30, ge=1, le=60, description="Maximum frames per second" ) include_metadata: bool = Field( default=True, description="Include metadata in stream" ) class StreamStatus(BaseModel): """Stream status model.""" is_active: bool = Field(..., description="Whether streaming is active") connected_clients: int = Field(..., description="Number of connected clients") streams: List[Dict[str, Any]] = Field(..., description="Active streams") uptime_seconds: float = Field(..., description="Stream uptime in seconds") # WebSocket endpoints @router.websocket("/pose") async def websocket_pose_stream( websocket: WebSocket, zone_ids: Optional[str] = Query(None, description="Comma-separated zone IDs"), min_confidence: float = Query(0.5, ge=0.0, le=1.0), max_fps: int = Query(30, ge=1, le=60), token: Optional[str] = Query(None, description="Authentication token") ): """WebSocket endpoint for real-time pose data streaming.""" client_id = None try: # Accept WebSocket connection await websocket.accept() # Check authentication if enabled from src.config.settings import get_settings settings = get_settings() if settings.enable_authentication and not token: await websocket.send_json({ "type": "error", "message": "Authentication token required" }) await websocket.close(code=1008) return # Parse zone IDs zone_list = None if zone_ids: zone_list = [zone.strip() for zone in zone_ids.split(",") if zone.strip()] # Register client with connection manager client_id = await connection_manager.connect( websocket=websocket, stream_type="pose", zone_ids=zone_list, min_confidence=min_confidence, max_fps=max_fps ) logger.info(f"WebSocket client {client_id} connected for pose streaming") # Send initial connection confirmation await websocket.send_json({ "type": "connection_established", "client_id": client_id, "timestamp": datetime.utcnow().isoformat(), "config": { "zone_ids": zone_list, "min_confidence": min_confidence, "max_fps": max_fps } }) # Keep connection alive and handle incoming messages while True: try: # Wait for client messages (ping, config updates, etc.) message = await websocket.receive_text() data = json.loads(message) await handle_websocket_message(client_id, data, websocket) except WebSocketDisconnect: break except json.JSONDecodeError: await websocket.send_json({ "type": "error", "message": "Invalid JSON format" }) except Exception as e: logger.error(f"Error handling WebSocket message: {e}") await websocket.send_json({ "type": "error", "message": "Internal server error" }) except WebSocketDisconnect: logger.info(f"WebSocket client {client_id} disconnected") except Exception as e: logger.error(f"WebSocket error: {e}") finally: if client_id: await connection_manager.disconnect(client_id) @router.websocket("/events") async def websocket_events_stream( websocket: WebSocket, event_types: Optional[str] = Query(None, description="Comma-separated event types"), zone_ids: Optional[str] = Query(None, description="Comma-separated zone IDs"), token: Optional[str] = Query(None, description="Authentication token") ): """WebSocket endpoint for real-time event streaming.""" client_id = None try: await websocket.accept() # Check authentication if enabled from src.config.settings import get_settings settings = get_settings() if settings.enable_authentication and not token: await websocket.send_json({ "type": "error", "message": "Authentication token required" }) await websocket.close(code=1008) return # Parse parameters event_list = None if event_types: event_list = [event.strip() for event in event_types.split(",") if event.strip()] zone_list = None if zone_ids: zone_list = [zone.strip() for zone in zone_ids.split(",") if zone.strip()] # Register client client_id = await connection_manager.connect( websocket=websocket, stream_type="events", zone_ids=zone_list, event_types=event_list ) logger.info(f"WebSocket client {client_id} connected for event streaming") # Send confirmation await websocket.send_json({ "type": "connection_established", "client_id": client_id, "timestamp": datetime.utcnow().isoformat(), "config": { "event_types": event_list, "zone_ids": zone_list } }) # Handle messages while True: try: message = await websocket.receive_text() data = json.loads(message) await handle_websocket_message(client_id, data, websocket) except WebSocketDisconnect: break except Exception as e: logger.error(f"Error in events WebSocket: {e}") except WebSocketDisconnect: logger.info(f"Events WebSocket client {client_id} disconnected") except Exception as e: logger.error(f"Events WebSocket error: {e}") finally: if client_id: await connection_manager.disconnect(client_id) async def handle_websocket_message(client_id: str, data: Dict[str, Any], websocket: WebSocket): """Handle incoming WebSocket messages.""" message_type = data.get("type") if message_type == "ping": await websocket.send_json({ "type": "pong", "timestamp": datetime.utcnow().isoformat() }) elif message_type == "update_config": # Update client configuration config = data.get("config", {}) await connection_manager.update_client_config(client_id, config) await websocket.send_json({ "type": "config_updated", "timestamp": datetime.utcnow().isoformat(), "config": config }) elif message_type == "get_status": # Send current status status = await connection_manager.get_client_status(client_id) await websocket.send_json({ "type": "status", "timestamp": datetime.utcnow().isoformat(), "status": status }) else: await websocket.send_json({ "type": "error", "message": f"Unknown message type: {message_type}" }) # HTTP endpoints for stream management @router.get("/status", response_model=StreamStatus) async def get_stream_status( stream_service: StreamService = Depends(get_stream_service) ): """Get current streaming status.""" try: status = await stream_service.get_status() connections = await connection_manager.get_connection_stats() # Calculate uptime (simplified for now) uptime_seconds = 0.0 if status.get("running", False): uptime_seconds = 3600.0 # Default 1 hour for demo return StreamStatus( is_active=status.get("running", False), connected_clients=connections.get("total_clients", status["connections"]["active"]), streams=[{ "type": "pose_stream", "active": status.get("running", False), "buffer_size": status["buffers"]["pose_buffer_size"] }], uptime_seconds=uptime_seconds ) except Exception as e: logger.error(f"Error getting stream status: {e}") raise HTTPException( status_code=500, detail=f"Failed to get stream status: {str(e)}" ) @router.post("/start") async def start_streaming( stream_service: StreamService = Depends(get_stream_service), current_user: Dict = Depends(require_auth) ): """Start the streaming service.""" try: logger.info(f"Starting streaming service by user: {current_user['id']}") if await stream_service.is_active(): return JSONResponse( status_code=200, content={"message": "Streaming service is already active"} ) await stream_service.start() return { "message": "Streaming service started successfully", "timestamp": datetime.utcnow().isoformat() } except Exception as e: logger.error(f"Error starting streaming: {e}") raise HTTPException( status_code=500, detail=f"Failed to start streaming: {str(e)}" ) @router.post("/stop") async def stop_streaming( stream_service: StreamService = Depends(get_stream_service), current_user: Dict = Depends(require_auth) ): """Stop the streaming service.""" try: logger.info(f"Stopping streaming service by user: {current_user['id']}") await stream_service.stop() await connection_manager.disconnect_all() return { "message": "Streaming service stopped successfully", "timestamp": datetime.utcnow().isoformat() } except Exception as e: logger.error(f"Error stopping streaming: {e}") raise HTTPException( status_code=500, detail=f"Failed to stop streaming: {str(e)}" ) @router.get("/clients") async def get_connected_clients( current_user: Dict = Depends(require_auth) ): """Get list of connected WebSocket clients.""" try: clients = await connection_manager.get_connected_clients() return { "total_clients": len(clients), "clients": clients, "timestamp": datetime.utcnow().isoformat() } except Exception as e: logger.error(f"Error getting connected clients: {e}") raise HTTPException( status_code=500, detail=f"Failed to get connected clients: {str(e)}" ) @router.delete("/clients/{client_id}") async def disconnect_client( client_id: str, current_user: Dict = Depends(require_auth) ): """Disconnect a specific WebSocket client.""" try: logger.info(f"Disconnecting client {client_id} by user: {current_user['id']}") success = await connection_manager.disconnect(client_id) if not success: raise HTTPException( status_code=404, detail=f"Client {client_id} not found" ) return { "message": f"Client {client_id} disconnected successfully", "timestamp": datetime.utcnow().isoformat() } except HTTPException: raise except Exception as e: logger.error(f"Error disconnecting client: {e}") raise HTTPException( status_code=500, detail=f"Failed to disconnect client: {str(e)}" ) @router.post("/broadcast") async def broadcast_message( message: Dict[str, Any], stream_type: Optional[str] = Query(None, description="Target stream type"), zone_ids: Optional[List[str]] = Query(None, description="Target zone IDs"), current_user: Dict = Depends(require_auth) ): """Broadcast a message to connected WebSocket clients.""" try: logger.info(f"Broadcasting message by user: {current_user['id']}") # Add metadata to message broadcast_data = { **message, "broadcast_timestamp": datetime.utcnow().isoformat(), "sender": current_user["id"] } # Broadcast to matching clients sent_count = await connection_manager.broadcast( data=broadcast_data, stream_type=stream_type, zone_ids=zone_ids ) return { "message": "Broadcast sent successfully", "recipients": sent_count, "timestamp": datetime.utcnow().isoformat() } except Exception as e: logger.error(f"Error broadcasting message: {e}") raise HTTPException( status_code=500, detail=f"Failed to broadcast message: {str(e)}" ) @router.get("/metrics") async def get_streaming_metrics(): """Get streaming performance metrics.""" try: metrics = await connection_manager.get_metrics() return { "metrics": metrics, "timestamp": datetime.utcnow().isoformat() } except Exception as e: logger.error(f"Error getting streaming metrics: {e}") raise HTTPException( status_code=500, detail=f"Failed to get streaming metrics: {str(e)}" )