import asyncio
from collections import defaultdict

from fastapi import WebSocket

from led_platform.domain import CommandEnvelope


class WebSocketHub:
    def __init__(self) -> None:
        self._lock = asyncio.Lock()
        self._groups: dict[str, set[WebSocket]] = defaultdict(set)

    async def connect(self, group: str, websocket: WebSocket) -> None:
        await websocket.accept()
        async with self._lock:
            self._groups[group].add(websocket)

    async def disconnect(self, group: str, websocket: WebSocket) -> None:
        async with self._lock:
            if group in self._groups:
                self._groups[group].discard(websocket)
                if not self._groups[group]:
                    self._groups.pop(group, None)

    async def send(self, websocket: WebSocket, message: CommandEnvelope) -> None:
        await websocket.send_json(message.model_dump(mode="json"))

    async def broadcast(self, message: CommandEnvelope, *groups: str) -> None:
        async with self._lock:
            targets: list[tuple[str, WebSocket]] = []
            for group in groups:
                targets.extend((group, websocket) for websocket in self._groups.get(group, set()))

        stale: list[tuple[str, WebSocket]] = []
        for group, websocket in targets:
            try:
                await self.send(websocket, message)
            except RuntimeError:
                stale.append((group, websocket))

        if stale:
            async with self._lock:
                for group, websocket in stale:
                    self._groups[group].discard(websocket)

    async def counts(self) -> dict[str, int]:
        async with self._lock:
            return {group: len(websockets) for group, websockets in self._groups.items()}