from fastapi import APIRouter, WebSocket

from app.websocket.web_socket_server import room_manager


web_socket = APIRouter()


@web_socket.websocket("/{room}")
async def websocket_room(websocket: WebSocket, room: str):
    """
    websocket 房间管理
    :param websocket:
    :param room:
    :return:
    """
    await websocket.accept()
    await room_manager.add_to_room(room, websocket)
    try:
        while True:
            data = await websocket.receive_text()
            await room_manager.broadcast_to_room(room, data, exclude_websocket=websocket)
    except Exception as e:
        print(f"连接关闭: {e}")
    finally:
        await room_manager.remove_from_room(room, websocket)
        await websocket.close()