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

"""
Gestor de conexiones WebSocket.

Maneja las conexiones de clientes WebSocket para
actualizaciones en tiempo real.
"""

import asyncio
import json
from datetime import datetime, timezone
from typing import Dict, List, Optional, Set

from fastapi import WebSocket


class ConnectionManager:
    """
    Gestor de conexiones WebSocket.

    Mantiene un registro de conexiones activas y permite
    enviar mensajes a clientes específicos o a todos.
    """

    def __init__(self):
        """Inicializa el gestor de conexiones."""
        # Conexiones activas por tipo de suscripción
        self._connections: Dict[str, Set[WebSocket]] = {
            "ubicaciones": set(),
            "alertas": set(),
            "vehiculos": set(),
        }
        # Conexiones por usuario
        self._user_connections: Dict[int, Set[WebSocket]] = {}
        # Suscripciones a vehículos específicos
        self._vehicle_subscriptions: Dict[int, Set[WebSocket]] = {}
        # Lock para operaciones thread-safe
        self._lock = asyncio.Lock()

    async def connect(
        self,
        websocket: WebSocket,
        channel: str = "ubicaciones",
        user_id: Optional[int] = None,
    ) -> None:
        """
        Acepta una nueva conexión WebSocket.

        Args:
            websocket: Conexión WebSocket.
            channel: Canal de suscripción.
            user_id: ID del usuario (opcional).
        """
        await websocket.accept()

        async with self._lock:
            # Agregar a conexiones del canal
            if channel in self._connections:
                self._connections[channel].add(websocket)

            # Agregar a conexiones del usuario
            if user_id:
                if user_id not in self._user_connections:
                    self._user_connections[user_id] = set()
                self._user_connections[user_id].add(websocket)

    async def disconnect(
        self,
        websocket: WebSocket,
        channel: str = "ubicaciones",
        user_id: Optional[int] = None,
    ) -> None:
        """
        Desconecta un WebSocket.

        Args:
            websocket: Conexión WebSocket.
            channel: Canal de suscripción.
            user_id: ID del usuario (opcional).
        """
        async with self._lock:
            # Remover de conexiones del canal
            if channel in self._connections:
                self._connections[channel].discard(websocket)

            # Remover de conexiones del usuario
            if user_id and user_id in self._user_connections:
                self._user_connections[user_id].discard(websocket)
                if not self._user_connections[user_id]:
                    del self._user_connections[user_id]

            # Remover de suscripciones de vehículos
            for vehicle_id in list(self._vehicle_subscriptions.keys()):
                self._vehicle_subscriptions[vehicle_id].discard(websocket)
                if not self._vehicle_subscriptions[vehicle_id]:
                    del self._vehicle_subscriptions[vehicle_id]

    async def subscribe_vehicle(
        self,
        websocket: WebSocket,
        vehicle_id: int,
    ) -> None:
        """
        Suscribe un WebSocket a actualizaciones de un vehículo específico.

        Args:
            websocket: Conexión WebSocket.
            vehicle_id: ID del vehículo.
        """
        async with self._lock:
            if vehicle_id not in self._vehicle_subscriptions:
                self._vehicle_subscriptions[vehicle_id] = set()
            self._vehicle_subscriptions[vehicle_id].add(websocket)

    async def unsubscribe_vehicle(
        self,
        websocket: WebSocket,
        vehicle_id: int,
    ) -> None:
        """
        Desuscribe un WebSocket de un vehículo.

        Args:
            websocket: Conexión WebSocket.
            vehicle_id: ID del vehículo.
        """
        async with self._lock:
            if vehicle_id in self._vehicle_subscriptions:
                self._vehicle_subscriptions[vehicle_id].discard(websocket)

    async def broadcast(
        self,
        message: dict,
        channel: str = "ubicaciones",
    ) -> None:
        """
        Envía un mensaje a todos los clientes de un canal.

        Args:
            message: Mensaje a enviar.
            channel: Canal de destino.
        """
        if channel not in self._connections:
            return

        message_json = json.dumps(message, default=str)
        disconnected = []

        for websocket in self._connections[channel]:
            try:
                await websocket.send_text(message_json)
            except Exception:
                disconnected.append(websocket)

        # Limpiar conexiones desconectadas
        for ws in disconnected:
            await self.disconnect(ws, channel)

    async def broadcast_vehicle_update(
        self,
        vehicle_id: int,
        data: dict,
    ) -> None:
        """
        Envía actualización a suscriptores de un vehículo específico.

        Args:
            vehicle_id: ID del vehículo.
            data: Datos a enviar.
        """
        message = {
            "type": "vehicle_update",
            "vehicle_id": vehicle_id,
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "data": data,
        }
        message_json = json.dumps(message, default=str)

        # Enviar a suscriptores del vehículo
        if vehicle_id in self._vehicle_subscriptions:
            disconnected = []
            for websocket in self._vehicle_subscriptions[vehicle_id]:
                try:
                    await websocket.send_text(message_json)
                except Exception:
                    disconnected.append(websocket)

            for ws in disconnected:
                await self.unsubscribe_vehicle(ws, vehicle_id)

        # También enviar al canal general de ubicaciones
        await self.broadcast(message, "ubicaciones")

    async def send_to_user(
        self,
        user_id: int,
        message: dict,
    ) -> None:
        """
        Envía un mensaje a todas las conexiones de un usuario.

        Args:
            user_id: ID del usuario.
            message: Mensaje a enviar.
        """
        if user_id not in self._user_connections:
            return

        message_json = json.dumps(message, default=str)
        disconnected = []

        for websocket in self._user_connections[user_id]:
            try:
                await websocket.send_text(message_json)
            except Exception:
                disconnected.append(websocket)

        # Limpiar conexiones desconectadas
        for ws in disconnected:
            await self.disconnect(ws, user_id=user_id)

    async def send_alert(
        self,
        alert_data: dict,
    ) -> None:
        """
        Envía una alerta a todos los clientes suscritos.

        Args:
            alert_data: Datos de la alerta.
        """
        message = {
            "type": "alert",
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "data": alert_data,
        }
        await self.broadcast(message, "alertas")

    def get_connection_count(self) -> dict:
        """
        Obtiene el conteo de conexiones activas.

        Returns:
            Dict con conteo por canal.
        """
        return {
            channel: len(connections)
            for channel, connections in self._connections.items()
        }

    def get_vehicle_subscribers(self, vehicle_id: int) -> int:
        """
        Obtiene el número de suscriptores de un vehículo.

        Args:
            vehicle_id: ID del vehículo.

        Returns:
            Número de suscriptores.
        """
        if vehicle_id in self._vehicle_subscriptions:
            return len(self._vehicle_subscriptions[vehicle_id])
        return 0


# Instancia global del gestor de conexiones
manager = ConnectionManager()