# NIKAME GENERATED — DO NOT EDIT DIRECTLY
from typing import Dict, List, Set
from fastapi import WebSocket

class ConnectionManager:
    def __init__(self):
        # Room Name -> Set of WebSockets
        self.rooms: Dict[str, Set[WebSocket]] = {}

    async def connect(self, websocket: WebSocket, room_name: str):
        await websocket.accept()
        if room_name not in self.rooms:
            self.rooms[room_name] = set()
        self.rooms[room_name].add(websocket)

    def disconnect(self, websocket: WebSocket, room_name: str):
        if room_name in self.rooms:
            self.rooms[room_name].remove(websocket)
            if not self.rooms[room_name]:
                del self.rooms[room_name]

    async def broadcast(self, message: str, room_name: str):
        if room_name in self.rooms:
            for connection in self.rooms[room_name]:
                await connection.send_text(message)

manager = ConnectionManager()
