from aiohttp import web, WSMsgType
from typing import Dict, Set

from .util import wrap_async
from ..web_agent import WebAgent, LlmState

class LlmWebSocket:
    """
    WebSocket handler for LLM state changes.
    Broadcasts state updates to all connected clients.
    """
    
    def __init__(self, agent: WebAgent):
        self._agent = agent
        self._clients: Set[web.WebSocketResponse] = set()
        self._agent.add_llm_change_handler(wrap_async(self._handle_state_change))
        
    async def _broadcast_message(self, message: Dict):
        """Broadcast message to all connected clients."""
        disconnected = set()
        for ws in self._clients:
            try:
                await ws.send_json(message)
            except ConnectionResetError:
                disconnected.add(ws)
        self._clients -= disconnected

    async def _handle_state_change(self, llm_name: str, new_state: LlmState):
        """Handle state changes from the WebAgent."""
        await self._broadcast_message({
            "llm": llm_name,
            "state": new_state.name
        })

    async def handle_connection(self, request: web.Request) -> web.WebSocketResponse:
        """Handle new WebSocket connections."""
        ws = web.WebSocketResponse(heartbeat=30)
        await ws.prepare(request)
        
        try:
            # Send initial states for all LLMs
            states = self._agent.llms
            for llm_name, state in states.items():
                await ws.send_json({
                    "llm": llm_name,
                    "state": state.name
                })
        
            self._clients.add(ws)
                
            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