import json
from collections import defaultdict

from fastapi import WebSocket


class Hub:
    def __init__(self) -> None:
        self.rooms: dict[str, list[WebSocket]] = defaultdict(list)

    async def join(self, session_id: str, ws: WebSocket) -> None:
        await ws.accept()
        self.rooms[session_id].append(ws)

    def leave(self, session_id: str, ws: WebSocket) -> None:
        conns = self.rooms.get(session_id, [])
        if ws in conns:
            conns.remove(ws)

    async def broadcast(self, session_id: str, payload: dict) -> None:
        dead = []
        for ws in list(self.rooms.get(session_id, [])):
            try:
                await ws.send_text(json.dumps(payload))
            except Exception:
                dead.append(ws)
        for ws in dead:
            self.leave(session_id, ws)


hub = Hub()