"""WebSocket менеджер для real-time обновлений статуса задач."""

import json
import logging
from typing import Any

from fastapi import WebSocket

logger = logging.getLogger(__name__)


class ConnectionManager:
    """Менеджер WebSocket соединений с поддержкой подписки на задачи."""

    def __init__(self) -> None:
        # task_id -> список активных соединений
        self._connections: dict[str, list[WebSocket]] = {}

    async def connect(self, task_id: str, websocket: WebSocket) -> None:
        """Принять WebSocket соединение и зарегистрировать его для задачи."""
        await websocket.accept()
        if task_id not in self._connections:
            self._connections[task_id] = []
        self._connections[task_id].append(websocket)
        logger.info(f"WebSocket подключён к задаче {task_id!r}")

    def disconnect(self, task_id: str, websocket: WebSocket) -> None:
        """Удалить соединение из реестра."""
        if task_id in self._connections:
            try:
                self._connections[task_id].remove(websocket)
            except ValueError:
                pass
            if not self._connections[task_id]:
                del self._connections[task_id]
        logger.info(f"WebSocket отключён от задачи {task_id!r}")

    async def send_task_update(self, task_id: str, data: dict[str, Any]) -> None:
        """
        Отправить обновление статуса задачи всем подключённым клиентам.

        Args:
            task_id: ID задачи
            data: Словарь с обновлением (status, queue_position, result и т.д.)
        """
        connections = self._connections.get(task_id, [])
        if not connections:
            return

        message = json.dumps(data, ensure_ascii=False, default=str)
        dead_connections = []

        for websocket in connections:
            try:
                await websocket.send_text(message)
            except Exception as e:
                logger.warning(f"Ошибка отправки WebSocket сообщения: {e}")
                dead_connections.append(websocket)

        # Очистить мёртвые соединения
        for ws in dead_connections:
            self.disconnect(task_id, ws)

    async def broadcast(self, data: dict[str, Any]) -> None:
        """Отправить сообщение всем подключённым клиентам."""
        message = json.dumps(data, ensure_ascii=False, default=str)
        all_dead = []

        for task_id, connections in self._connections.items():
            for websocket in connections:
                try:
                    await websocket.send_text(message)
                except Exception:
                    all_dead.append((task_id, websocket))

        for task_id, ws in all_dead:
            self.disconnect(task_id, ws)

    @property
    def active_connections_count(self) -> int:
        """Количество активных WebSocket соединений."""
        return sum(len(conns) for conns in self._connections.values())


# Глобальный экземпляр менеджера
ws_manager = ConnectionManager()