"""
WebSocket client for real-time feeds.

Provides clients for:
- Pulse feed (/ws/v1/pulse) — live enriched headlines
- Smart columns (/ws/v1/smart-columns) — real-time match delivery
- Enrichment feed (/ws/v1/enrichment/feed) — bidirectional enrichment
"""
import asyncio
import json
import logging
from typing import Any, Callable, Dict, Optional

try:
    import websockets
    from websockets.client import WebSocketClientProtocol
except ImportError:
    websockets = None  # type: ignore[assignment]

logger = logging.getLogger("kronical.websocket")


class KronicalWebSocket:
    """
    WebSocket client with automatic reconnection and heartbeat handling.

    Usage:
        ws = KronicalWebSocket(
            api_key="kron_hf_...",
            base_url="wss://api.kronical.ai",
        )

        # Pulse feed
        async def on_headline(msg):
            print(msg["text"])

        await ws.pulse(on_message=on_headline)

        # Enrichment feed
        await ws.connect_enrichment()
        result = await ws.enrich("Fed raises rates by 25bps")
    """

    def __init__(
        self,
        api_key: str,
        base_url: str,
        reconnect_delay: float = 5.0,
        max_reconnect_delay: float = 60.0,
    ):
        if websockets is None:
            raise ImportError(
                "websockets package required for WebSocket support. "
                "Install with: pip install websockets"
            )
        self._api_key = api_key
        self._base_url = base_url.rstrip("/")
        self._reconnect_delay = reconnect_delay
        self._max_reconnect_delay = max_reconnect_delay
        self._ws: Optional[WebSocketClientProtocol] = None
        self._running = False

    def _ws_url(self, path: str) -> str:
        """Build WebSocket URL with auth token as query param."""
        return f"{self._base_url}{path}?token={self._api_key}"

    async def pulse(
        self,
        on_message: Callable[[Dict[str, Any]], Any],
        *,
        filters: Optional[Dict[str, str]] = None,
    ) -> None:
        """
        Connect to the pulse WebSocket and stream headlines.

        On connection, the server sends a ``connected`` message that includes
        a ``pulse_columns`` list with the current column configuration.

        Each headline message includes a ``matched_columns`` list containing
        the IDs of pulse columns that matched the headline.

        Args:
            on_message: Callback for each message (connected, headline, etc.).
            filters: Optional filters (novelty, escalation, tone).

        Runs forever with automatic reconnection.
        """
        self._running = True
        delay = self._reconnect_delay

        while self._running:
            try:
                url = self._ws_url("/ws/v1/pulse")
                async with websockets.connect(url) as ws:
                    self._ws = ws
                    logger.info("Connected to pulse feed")
                    delay = self._reconnect_delay

                    # Send subscription filters
                    if filters:
                        await ws.send(json.dumps({
                            "type": "subscribe",
                            **filters,
                        }))

                    async for raw in ws:
                        msg = json.loads(raw)

                        # Handle heartbeat
                        if msg.get("type") == "ping":
                            await ws.send(json.dumps({"type": "pong"}))
                            continue

                        if msg.get("type") == "server_shutdown":
                            reconnect_ms = msg.get("reconnect_after_ms", 5000)
                            logger.info(f"Server shutting down, reconnecting in {reconnect_ms}ms")
                            await asyncio.sleep(reconnect_ms / 1000)
                            break

                        await _maybe_await(on_message, msg)

            except Exception as e:
                if not self._running:
                    break
                logger.warning(f"Pulse connection lost: {e}, reconnecting in {delay}s")
                await asyncio.sleep(delay)
                delay = min(delay * 2, self._max_reconnect_delay)
            finally:
                self._ws = None

    async def refresh_columns(self) -> None:
        """
        Send a refresh_columns message on the pulse WebSocket.

        Tells the server to re-evaluate pulse column configuration.
        Must be called while connected to the pulse feed.
        """
        if not self._ws:
            raise RuntimeError("Not connected — call pulse() first")
        await self._ws.send(json.dumps({"type": "refresh_columns"}))

    async def smart_columns(
        self,
        on_message: Callable[[Dict[str, Any]], Any],
    ) -> None:
        """
        Connect to the smart columns WebSocket for real-time matches.

        Runs forever with automatic reconnection.
        """
        self._running = True
        delay = self._reconnect_delay

        while self._running:
            try:
                url = self._ws_url("/ws/v1/smart-columns")
                async with websockets.connect(url) as ws:
                    logger.info("Connected to smart columns feed")
                    delay = self._reconnect_delay

                    async for raw in ws:
                        msg = json.loads(raw)
                        if msg.get("type") == "ping":
                            await ws.send(json.dumps({"type": "pong"}))
                            continue
                        if msg.get("type") == "server_shutdown":
                            await asyncio.sleep(msg.get("reconnect_after_ms", 5000) / 1000)
                            break
                        await _maybe_await(on_message, msg)

            except Exception as e:
                if not self._running:
                    break
                logger.warning(f"Smart columns connection lost: {e}, reconnecting in {delay}s")
                await asyncio.sleep(delay)
                delay = min(delay * 2, self._max_reconnect_delay)

    async def connect_enrichment(self) -> None:
        """Connect to the enrichment WebSocket for bidirectional enrichment."""
        url = self._ws_url("/ws/v1/enrichment/feed")
        self._ws = await websockets.connect(url)
        self._running = True
        logger.info("Connected to enrichment feed")

    async def enrich(self, headline_text: str, *, headline_id: Optional[str] = None, source: str = "sdk") -> Dict[str, Any]:
        """
        Send a headline for enrichment and wait for the result.

        Args:
            headline_text: Headline to enrich.
            headline_id: Optional client-side ID for correlation.
            source: Source identifier.

        Returns:
            Enrichment result dict.
        """
        if not self._ws:
            raise RuntimeError("Not connected — call connect_enrichment() first")

        msg = {
            "type": "headline",
            "id": headline_id or "",
            "text": headline_text,
            "source": source,
        }
        await self._ws.send(json.dumps(msg))

        # Wait for response (handle heartbeats in between)
        while True:
            raw = await self._ws.recv()
            response = json.loads(raw)
            if response.get("type") == "ping":
                await self._ws.send(json.dumps({"type": "pong"}))
                continue
            return response

    async def close(self) -> None:
        """Disconnect all WebSocket connections."""
        self._running = False
        if self._ws:
            await self._ws.close()
            self._ws = None

    async def __aenter__(self):
        return self

    async def __aexit__(self, *args):
        await self.close()


async def _maybe_await(func: Callable, *args: Any) -> Any:
    """Call a function, awaiting if it returns a coroutine."""
    result = func(*args)
    if asyncio.iscoroutine(result):
        return await result
    return result
