from aiohttp import web

from ..web_agent import WebAgent
from .util import wrap_async


class ResponseWebSocket:
    def __init__(self, agent: WebAgent):
        self._agent = agent
        self._clients = set()
        self._agent._response_buffer.add_observer(wrap_async(self._handle_buffer_change))
    
    async def _handle_buffer_change(self, buffer):
        await self._broadcast_message({
            "response": buffer,
        })
    
    async def _broadcast_message(self, message):
        disconnected = set()
        for ws in self._clients:
            try:
                await ws.send_json(message)
            except ConnectionResetError:
                disconnected.add(ws)
        self._clients -= disconnected
    
    async def handle_connection(self, request: web.Request) -> web.WebSocketResponse:
        ws = web.WebSocketResponse(heartbeat=30)
        await ws.prepare(request)
        
        self._clients.add(ws)
        
        try:
            await ws.send_json({
                "response": self._agent.response_buffer.get_text(),
            })
                
            async for msg in ws:
                if msg.type == WSMsgType.ERROR:
                    print(f"WebSocket connection closed with error: {ws.exception()}")
        finally:
            self._clients.remove(ws)
            
        return ws