"""WebSocket manager for broadcasting real-time evaluation progress."""

import json
from collections import defaultdict
from typing import Any

from fastapi import WebSocket


class ConnectionManager:
    """Manages WebSocket connections grouped by run_id."""

    def __init__(self):
        self._connections: dict[str, list[WebSocket]] = defaultdict(list)

    async def connect(self, run_id: str, websocket: WebSocket):
        await websocket.accept()
        self._connections[run_id].append(websocket)

    def disconnect(self, run_id: str, websocket: WebSocket):
        conns = self._connections.get(run_id, [])
        if websocket in conns:
            conns.remove(websocket)
        if not conns:
            self._connections.pop(run_id, None)

    async def broadcast(self, run_id: str, data: dict[str, Any]):
        conns = self._connections.get(run_id, [])
        message = json.dumps(data, ensure_ascii=False, default=str)
        dead = []
        for ws in conns:
            try:
                await ws.send_text(message)
            except Exception:
                dead.append(ws)
        for ws in dead:
            self.disconnect(run_id, ws)

    async def emit(self, run_id: str, event: str, data: dict[str, Any]) -> None:
        """Async progress-callback compatible with EvalEngine's awaitable callback."""
        await self.broadcast(run_id, {"event": event, **data})


ws_manager = ConnectionManager()