"""WebSocket router: live vote updates."""

from __future__ import annotations

import json
import uuid
from typing import Any

from fastapi import APIRouter, WebSocket, WebSocketDisconnect

router = APIRouter()


# ── Connection manager ──────────────────────────────────────────────────────


class ConnectionManager:
    """Manages active WebSocket connections grouped by vote session ID."""

    def __init__(self) -> None:
        # session_id -> list of connected websockets
        self._connections: dict[uuid.UUID, list[WebSocket]] = {}

    async def connect(self, websocket: WebSocket, session_id: uuid.UUID) -> None:
        """Accept a WebSocket connection and register it for a vote session."""
        await websocket.accept()
        if session_id not in self._connections:
            self._connections[session_id] = []
        self._connections[session_id].append(websocket)

    def disconnect(self, websocket: WebSocket, session_id: uuid.UUID) -> None:
        """Remove a WebSocket connection from the session group."""
        if session_id in self._connections:
            self._connections[session_id] = [
                ws for ws in self._connections[session_id] if ws is not websocket
            ]
            if not self._connections[session_id]:
                del self._connections[session_id]

    async def broadcast(self, session_id: uuid.UUID, data: dict[str, Any]) -> None:
        """Broadcast a message to all connections watching a given vote session."""
        if session_id not in self._connections:
            return

        message = json.dumps(data, default=str)
        dead: list[WebSocket] = []

        for ws in self._connections[session_id]:
            try:
                await ws.send_text(message)
            except Exception:
                dead.append(ws)

        # Clean up dead connections
        for ws in dead:
            self.disconnect(ws, session_id)


manager = ConnectionManager()


# ── WebSocket endpoint ──────────────────────────────────────────────────────


@router.websocket("/live")
async def live_updates(websocket: WebSocket) -> None:
    """WebSocket endpoint for live vote session updates.

    The client connects and sends a JSON message with the session_id
    they want to subscribe to:

        { "action": "subscribe", "session_id": "<uuid>" }

    The server will then push vote update events to the client:

        { "event": "vote_update", "session_id": "...", "votes_for": N, "votes_against": N, "votes_total": N }
        { "event": "session_closed", "session_id": "...", "result": "adopted|rejected" }

    The client can also unsubscribe:

        { "action": "unsubscribe", "session_id": "<uuid>" }
    """
    await websocket.accept()
    subscribed_sessions: set[uuid.UUID] = set()

    try:
        while True:
            raw = await websocket.receive_text()

            try:
                data = json.loads(raw)
            except json.JSONDecodeError:
                await websocket.send_text(json.dumps({"error": "JSON invalide"}))
                continue

            action = data.get("action")
            session_id_str = data.get("session_id")

            if not action or not session_id_str:
                await websocket.send_text(
                    json.dumps({"error": "Champs 'action' et 'session_id' requis"})
                )
                continue

            try:
                session_id = uuid.UUID(session_id_str)
            except ValueError:
                await websocket.send_text(json.dumps({"error": "session_id invalide"}))
                continue

            if action == "subscribe":
                if session_id not in subscribed_sessions:
                    # Register this websocket in the manager for this session
                    if session_id not in manager._connections:
                        manager._connections[session_id] = []
                    manager._connections[session_id].append(websocket)
                    subscribed_sessions.add(session_id)

                await websocket.send_text(
                    json.dumps({"event": "subscribed", "session_id": str(session_id)})
                )

            elif action == "unsubscribe":
                if session_id in subscribed_sessions:
                    manager.disconnect(websocket, session_id)
                    subscribed_sessions.discard(session_id)

                await websocket.send_text(
                    json.dumps({"event": "unsubscribed", "session_id": str(session_id)})
                )

            else:
                await websocket.send_text(
                    json.dumps({"error": f"Action inconnue: {action}"})
                )

    except WebSocketDisconnect:
        # Clean up all subscriptions for this client
        for session_id in subscribed_sessions:
            manager.disconnect(websocket, session_id)