from fastapi import WebSocket, WebSocketDisconnect
from . import ws_router
from ..bot.chatbot import ChatBot


@ws_router.websocket("/ws/{client_id}")
async def websocket_endpoint(websocket: WebSocket, client_id: str):
    await websocket.accept()
    chatbot = ChatBot()
    try:
        while True:
            data = await websocket.receive_json()
            answer = await chatbot.ask(data)
            await websocket.send_json(answer)
    except WebSocketDisconnect:
        pass