"""
WebSocket channel adapter.
Created: 2026-02-02
Changes:
  - 2026-02-05: Fixed system_event format - send flat structure for frontend
"""

import logging
from typing import Any

from fastapi import WebSocket

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

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."""
        # Send flat structure: {type, event_type, data}
        # Frontend expects event_type and data at top level
        payload = {"type": "system_event", "event_type": event.event_type, "data": event.data}

        # Broadcast to all connected clients
        for ws in self._connections.values():
            try:
                await ws.send_json(payload)
            except Exception:
                pass

    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
