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

from .util import wrap_async
from ..chat_io_buffer import ChatIOBuffer

class ChatWebSocket:
    
    def __init__(self, io_buffer: ChatIOBuffer):
        self._io_buffer = io_buffer
        self._clients: Set[web.WebSocketResponse] = set()
        self._io_buffer.add_change_handler(wrap_async(self._handle_io_buffer_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_io_buffer_change(self, new_messages, deleted_message_ids, last_read_timestamp):
        """Handle changes to the ChatIOBuffer."""
        await self._broadcast_message({
            "messages": [
                {
                    "id": message.id,
                    "content": message.content,
                    "message_type": message.message_type.value
                }
                for message in new_messages
            ],
            "deleted_message_ids": deleted_message_ids,
            "last_read_timestamp": last_read_timestamp
        })

    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 state
            await ws.send_json({
                "messages": [
                    {
                        "id": message.id,
                        "content": message.content,
                        "message_type": message.message_type.value
                    }
                    for message in self._io_buffer.messages
                ],
                "last_read_timestamp": self._io_buffer.last_read_timestamp
            })
            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