"""
WebSocket channel adapter.
Created: 2026-02-02
"""

from fastapi import WebSocket
from typing import Any
import json
import logging

from pocketclaw.bus.adapters import BaseChannelAdapter
from pocketclaw.bus.events import Channel, InboundMessage, OutboundMessage, SystemEvent
from pocketclaw.bus.queue import MessageBus
from dataclasses import asdict

logger = logging.getLogger(__name__)


class WebSocketAdapter(BaseChannelAdapter):
    """
    WebSocket channel adapter.

    Manages multiple WebSocket connections and routes messages appropriately.
    """

    def __init__(self):
        super().__init__()
        self._connections: dict[str, WebSocket] = {}  # chat_id -> WebSocket

    @property
    def channel(self) -> Channel:
        return Channel.WEBSOCKET

    async def start(self, bus: MessageBus) -> None:
        """Start adapter and subscribe to both outbound and system events."""
        await super().start(bus)
        # Subscribe to system events (thinking, tool usage, etc.)
        bus.subscribe_system(self.on_system_event)
        logger.info("🔌 WebSocket Adapter subscribed to System Events")

    async def on_system_event(self, event: SystemEvent) -> None:
        """Handle system event by broadcasting to all clients."""
        # Convert dataclass to dict
        event_dict = asdict(event)
        
        # Broadcast to all connected clients
        # In multi-user system, we'd filter by session_key/chat_id if available in data
        # For now, broadcast to all (dashboard is single view)
        await self.broadcast(event_dict, msg_type="system_event")

    async def register_connection(self, websocket: WebSocket, chat_id: str) -> None:
        """Register a new WebSocket connection."""
        # Assume connection is already accepted by the handler
        self._connections[chat_id] = websocket
        logger.info(f"🔌 WebSocket connected: {chat_id}")

    async def unregister_connection(self, chat_id: str) -> None:
        """Unregister a WebSocket connection."""
        self._connections.pop(chat_id, None)
        logger.info(f"🔌 WebSocket disconnected: {chat_id}")

    async def handle_message(self, chat_id: str, data: dict[str, Any]) -> None:
        """Handle incoming WebSocket message."""
        action = data.get("action", "chat")

        if action == "chat":
            message = InboundMessage(
                channel=Channel.WEBSOCKET,
                sender_id=chat_id,
                chat_id=chat_id,
                content=data.get("message", ""),
                metadata=data,
            )

            
            # Send stream_start to frontend to initialize the response UI
            ws = self._connections.get(chat_id)
            if ws:
                try:
                    await ws.send_json({"type": "stream_start"})
                except Exception:
                    pass
            
            await self._publish_inbound(message)
        # Other actions (settings, tools) handled separately

    async def send(self, message: OutboundMessage) -> None:
        """Send message to WebSocket client."""
        ws = self._connections.get(message.chat_id)
        if not ws:
            # Broadcast to all if no specific chat_id
            for ws in self._connections.values():
                await self._send_to_socket(ws, message)
        else:
            await self._send_to_socket(ws, message)

    async def _send_to_socket(self, ws: WebSocket, message: OutboundMessage) -> None:
        """Send to a specific WebSocket."""
        try:

            if message.is_stream_end:
                await ws.send_json({"type": "stream_end"})
                return

            # Legacy frontend expects 'message' for chunks too
            msg_type = "message" 
            
            # If it's the very first chunk, we might want to send stream_start? 
            # But we don't track state easily here. 
            # Let's hope frontend doesn't strictly need stream_start if it's already waiting.
            # (Legacy dashboard.py sent stream_start before headers)
            
            await ws.send_json({
                "type": msg_type,
                "content": message.content,
                "metadata": message.metadata,
            })
        except Exception:
            pass  # Connection closed

            pass  # Connection closed

    async def broadcast(self, content: Any, msg_type: str = "notification") -> None:
        """Broadcast to all connected clients."""
        for ws in self._connections.values():
            try:
                await ws.send_json({"type": msg_type, "content": content})
            except Exception:
                pass
