from fastapi import WebSocket


class SocketManager:
    def __init__(self):
        self.rooms = {}

    async def add_to_room(self, room: str, websocket: WebSocket):
        if room not in self.rooms:
            self.rooms[room] = []
        self.rooms[room].append(websocket)

    async def remove_from_room(self, room: str, websocket: WebSocket):
        if room in self.rooms:
            self.rooms[room].remove(websocket)
            if len(self.rooms[room]) == 0:
                del self.rooms[room]
                if room.startswith('detect_rtsp_'):
                    print()

    async def broadcast_to_room(self, room: str, message: str, exclude_websocket: WebSocket = None):
        if room in self.rooms:
            for ws in self.rooms[room]:
                if ws != exclude_websocket:
                    try:
                        await ws.send_text(message)
                    except:
                        await self.remove_from_room(room, ws)

    async def send_to_room(self, room: str, message: str):
        if room in self.rooms:
            for ws in self.rooms[room]:
                try:
                    await ws.send_text(message)
                except Exception as e:
                    print(e)

    async def send_stream_to_room(self, room: str, message):
        if room in self.rooms:
            for ws in self.rooms[room]:
                try:
                    await ws.send_bytes(message)
                except Exception as e:
                    print(e)


room_manager = SocketManager()