import asyncio
from typing import Dict, Optional
from fastapi import WebSocket


class ConnectionManager:
    def __init__(self):
        self.connections: Dict[str, WebSocket] = {}
        self.lock = asyncio.Lock()

    async def add_connection(self, device_id: str, websocket: WebSocket):
        async with self.lock:
            self.connections[device_id] = websocket

    async def remove_connection(self, device_id: str, websocket: Optional[WebSocket] = None):
        async with self.lock:
            if device_id in self.connections and (websocket is None or self.connections[device_id] is websocket):
                del self.connections[device_id]

    async def get_connection(self, device_id: str) -> Optional[WebSocket]:
        async with self.lock:
            return self.connections.get(device_id)

    async def is_connected(self, device_id: str) -> bool:
        websocket = await self.get_connection(device_id)
        return (
            websocket is not None
            and getattr(getattr(websocket, "client_state", None), "name", None) == "CONNECTED"
            and not getattr(websocket, "_closed", False)
        )

    async def get_all_connections(self) -> Dict[str, WebSocket]:
        async with self.lock:
            return dict(self.connections)


connection_manager = ConnectionManager()