# app/ws/connection_manager.py
from typing import Dict
from fastapi import WebSocket


class VisitorConnectionManager:
    """
    Keeps track of active visitor websocket connections.
    Keys are visitor_uids (strings), values are websocket objects.
    """

    def __init__(self):
        self.active_connections: dict[str, any] = {}

    def connect(self, visitor_uid: str, websocket):
        self.active_connections[visitor_uid] = websocket

    def disconnect(self, visitor_uid: str):
        if visitor_uid in self.active_connections:
            del self.active_connections[visitor_uid]

    def get(self, visitor_uid: str):
        return self.active_connections.get(visitor_uid)


class AdminConnectionManager:
    """
    Admins don’t have IDs (unless you want), we just broadcast to all.
    """
    def __init__(self):
        self.active_connections: list[WebSocket] = []

    def connect(self, websocket: WebSocket):
        self.active_connections.append(websocket)

    def disconnect(self, websocket: WebSocket):
        if websocket in self.active_connections:
            self.active_connections.remove(websocket)

    async def broadcast(self, message: dict):
        dead = []
        for ws in self.active_connections:
            try:
                await ws.send_json(message)
            except Exception:
                dead.append(ws)
        for ws in dead:
            self.disconnect(ws)


class AgentConnectionManager:
    def __init__(self):
        self.active_connections: Dict[int, WebSocket] = {}

    def connect(self, agent_id: int, websocket: WebSocket):
        self.active_connections[agent_id] = websocket

    def disconnect(self, agent_id: int):
        if agent_id in self.active_connections:
            del self.active_connections[agent_id]

    def get(self, agent_id: int):
        return self.active_connections.get(agent_id)

    async def broadcast(self, message: dict):
        dead = []
        for aid, ws in self.active_connections.items():
            try:
                await ws.send_json(message)
            except Exception:
                dead.append(aid)
        for aid in dead:
            self.disconnect(aid)


agent_manager = AgentConnectionManager()
admin_manager = AdminConnectionManager()
visitor_manager = VisitorConnectionManager()