ATLAS/backend/app/api/websocket/ubicaciones.py

"""
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)