# sanguo_api/ws.py
"""WS 连接池：task_id → 订阅者集合。简单广播，不做重连/心跳。"""
from fastapi import WebSocket


class ConnectionManager:
    def __init__(self):
        self._connections: dict[str, set[WebSocket]] = {}

    def connect(self, task_id: str, ws: WebSocket):
        self._connections.setdefault(task_id, set()).add(ws)

    def disconnect(self, task_id: str, ws: WebSocket):
        conns = self._connections.get(task_id)
        if conns:
            conns.discard(ws)
            if not conns:
                del self._connections[task_id]

    async def broadcast(self, task_id: str, msg: dict):
        for ws in list(self._connections.get(task_id, [])):
            try:
                await ws.send_json(msg)
            except Exception:
                self.disconnect(task_id, ws)


manager = ConnectionManager()