from fastapi import WebSocket
from fastapi import FastAPI
from starlette.websockets import WebSocketState

from utils.websocket_server import room_manager


def websocket_config(app: FastAPI):

    @app.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)