from fastapi import APIRouter, WebSocket
from starlette.websockets import WebSocketState

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:
        if websocket.client_state != WebSocketState.DISCONNECTED:
            await websocket.close(code=1000)
    finally:
        await room_manager.remove_from_room(room, websocket)
        if websocket.client_state != WebSocketState.DISCONNECTED:
            await websocket.close(code=1001)