from __future__ import annotations

import json
import logging
from collections import defaultdict

from fastapi import WebSocket

logger = logging.getLogger(__name__)


class ConnectionManager:
    """Per-tenant WebSocket registry.

    Maintains an in-process dict of tenant_id → set of active WebSocket
    connections.  Dead connections are pruned lazily during broadcast.

    NOTE: This implementation is node-local.  For multi-instance deployments,
    replace the backing store with Redis Pub/Sub.
    """

    def __init__(self) -> None:
        self._connections: dict[str, set[WebSocket]] = defaultdict(set)

    async def connect(self, websocket: WebSocket, tenant_id: str) -> None:
        await websocket.accept()
        self._connections[tenant_id].add(websocket)
        logger.debug("WS connected tenant=%s total=%d", tenant_id, len(self._connections[tenant_id]))

    def disconnect(self, websocket: WebSocket, tenant_id: str) -> None:
        self._connections[tenant_id].discard(websocket)

    async def broadcast_to_tenant(self, tenant_id: str, payload: dict) -> None:
        message = json.dumps(payload)
        dead: list[WebSocket] = []
        for ws in list(self._connections.get(tenant_id, [])):
            try:
                await ws.send_text(message)
            except Exception:
                dead.append(ws)
        for ws in dead:
            self._connections[tenant_id].discard(ws)


manager = ConnectionManager()