from aiohttp import web

from ..auto_approver import AutoApprover
from ..chat_io_buffer import ChatIOBuffer
from ..web_agent import WebAgent
from ..working_memory import WorkingMemory
from .auto_approver_websocket import AutoApproverWebSocket
from .chat_websocket import ChatWebSocket
from .memory_websocket import MemoryWebSocket
from .response_websocket import ResponseWebSocket
from .state_websocket import StateWebSocket

class Websockets:
    def __init__(self, app: web.Application, agent: WebAgent, io_buffer: ChatIOBuffer, auto_approver: AutoApprover, working_memory: WorkingMemory):
        self._auto_approver_ws = AutoApproverWebSocket(auto_approver)
        self._chat_ws = ChatWebSocket(io_buffer)
        self._memory_ws = MemoryWebSocket(working_memory)
        self._response_ws = ResponseWebSocket(agent)
        self._state_ws = StateWebSocket(agent)

        app.router.add_get("/ws/auto_approver", self._auto_approver_ws.handle_connection)
        app.router.add_get("/ws/chat", self._chat_ws.handle_connection)
        app.router.add_get("/ws/memory", self._memory_ws.handle_connection)
        app.router.add_get("/ws/response", self._response_ws.handle_connection)
        app.router.add_get("/ws/state", self._state_ws.handle_connection)