""" WebSocket endpoint para ubicaciones en tiempo real. """ import json from typing import Optional from fastapi import APIRouter, Depends, Query, WebSocket, WebSocketDisconnect from jose import JWTError from app.api.websocket.manager import manager from app.core.security import decode_token router = APIRouter() async def get_user_from_token(token: Optional[str]) -> Optional[int]: """ Obtiene el ID de usuario desde un token JWT. Args: token: Token JWT. Returns: ID del usuario o None. """ if not token: return None try: payload = decode_token(token) return int(payload.get("sub")) except (JWTError, ValueError): return None @router.websocket("/ws/ubicaciones") async def websocket_ubicaciones( websocket: WebSocket, token: Optional[str] = Query(None), ): """ WebSocket para recibir actualizaciones de ubicaciones. Permite suscribirse a: - Todas las ubicaciones de la flota - Vehículos específicos Args: websocket: Conexión WebSocket. token: Token JWT para autenticación (opcional). """ user_id = await get_user_from_token(token) await manager.connect(websocket, "ubicaciones", user_id) try: while True: # Recibir mensajes del cliente data = await websocket.receive_text() try: message = json.loads(data) action = message.get("action") if action == "subscribe_vehicle": # Suscribirse a un vehículo específico vehicle_id = message.get("vehicle_id") if vehicle_id: await manager.subscribe_vehicle(websocket, vehicle_id) await websocket.send_json({ "type": "subscribed", "vehicle_id": vehicle_id, }) elif action == "unsubscribe_vehicle": # Desuscribirse de un vehículo vehicle_id = message.get("vehicle_id") if vehicle_id: await manager.unsubscribe_vehicle(websocket, vehicle_id) await websocket.send_json({ "type": "unsubscribed", "vehicle_id": vehicle_id, }) elif action == "ping": # Responder ping para keepalive await websocket.send_json({"type": "pong"}) except json.JSONDecodeError: await websocket.send_json({ "type": "error", "message": "Invalid JSON", }) except WebSocketDisconnect: await manager.disconnect(websocket, "ubicaciones", user_id) @router.websocket("/ws/vehiculo/{vehiculo_id}") async def websocket_vehiculo( websocket: WebSocket, vehiculo_id: int, token: Optional[str] = Query(None), ): """ WebSocket para seguir un vehículo específico. Args: websocket: Conexión WebSocket. vehiculo_id: ID del vehículo a seguir. token: Token JWT para autenticación (opcional). """ user_id = await get_user_from_token(token) await manager.connect(websocket, "ubicaciones", user_id) await manager.subscribe_vehicle(websocket, vehiculo_id) # Enviar confirmación de suscripción await websocket.send_json({ "type": "connected", "vehicle_id": vehiculo_id, }) try: while True: data = await websocket.receive_text() try: message = json.loads(data) if message.get("action") == "ping": await websocket.send_json({"type": "pong"}) except json.JSONDecodeError: pass except WebSocketDisconnect: await manager.unsubscribe_vehicle(websocket, vehiculo_id) await manager.disconnect(websocket, "ubicaciones", user_id) @router.websocket("/ws/flota") async def websocket_flota( websocket: WebSocket, token: Optional[str] = Query(None), ): """ WebSocket para monitoreo de toda la flota. Recibe actualizaciones de todos los vehículos activos. Args: websocket: Conexión WebSocket. token: Token JWT para autenticación (opcional). """ user_id = await get_user_from_token(token) await manager.connect(websocket, "ubicaciones", user_id) # Enviar confirmación await websocket.send_json({ "type": "connected", "channel": "fleet", }) try: while True: data = await websocket.receive_text() try: message = json.loads(data) if message.get("action") == "ping": await websocket.send_json({"type": "pong"}) elif message.get("action") == "request_status": # Enviar estado actual de conexiones await websocket.send_json({ "type": "status", "connections": manager.get_connection_count(), }) except json.JSONDecodeError: pass except WebSocketDisconnect: await manager.disconnect(websocket, "ubicaciones", user_id)