"""WebSocket connection manager for real-time version room events."""

from __future__ import annotations

import json
import uuid
from typing import Any

from fastapi import WebSocket


class ConnectionManager:
    def __init__(self) -> None:
        # version_id -> list of active WebSocket connections
        self._rooms: dict[str, list[WebSocket]] = {}

    async def connect(self, version_id: uuid.UUID, websocket: WebSocket) -> None:
        await websocket.accept()
        key = str(version_id)
        self._rooms.setdefault(key, []).append(websocket)

    def disconnect(self, version_id: uuid.UUID, websocket: WebSocket) -> None:
        key = str(version_id)
        room = self._rooms.get(key, [])
        if websocket in room:
            room.remove(websocket)
        if not room:
            self._rooms.pop(key, None)

    async def broadcast(self, version_id: uuid.UUID, event: str, data: Any) -> None:
        key = str(version_id)
        payload = json.dumps({"event": event, "data": data})
        dead: list[WebSocket] = []
        for ws in self._rooms.get(key, []):
            try:
                await ws.send_text(payload)
            except Exception:
                dead.append(ws)
        for ws in dead:
            self.disconnect(version_id, ws)

    def room_size(self, version_id: uuid.UUID) -> int:
        return len(self._rooms.get(str(version_id), []))


manager = ConnectionManager()