"""autoplay_sdk.client — ConnectorClient for the Autoplay real-time event stream.

Receives structured UI actions and LLM session summaries from the Autoplay
event connector in real-time via SSE (Server-Sent Events).

Quickstart::

    from autoplay_sdk import ConnectorClient, ActionsPayload

    def on_actions(payload: ActionsPayload):
        for action in payload.actions:
            print(action.title, action.canonical_url)

    ConnectorClient(url="https://host/stream/my_product", token="uk_live_...") \\
        .on_actions(on_actions) \\
        .run()

Background usage (non-blocking)::

    client = ConnectorClient(url="...", token="...").on_actions(on_actions)
    client.run_in_background()
    # your application continues here
    ...
    client.stop()

Context manager (clean shutdown in scripts)::

    with ConnectorClient(url="...", token="...") as client:
        client.on_actions(on_actions).run_in_background()
        do_other_work()
    # stop() is called automatically on exit

High-volume / RAG usage
-----------------------
The client decouples the SSE receive loop from callback execution via an
internal ``queue.Queue``.  The receive thread reads events from the stream and
enqueues them; a separate worker thread drains the queue and calls your
callbacks.  Your callback can block (e.g. write to a vector store) without
ever stalling the stream reader.

If the queue fills up (``max_queue_size`` events are waiting), new events are
dropped, counted in ``dropped_count``, and passed to your ``on_drop`` callback
if one is registered.  Monitor ``dropped_count`` in production — a non-zero
value means your callback is slower than the incoming event rate.

Lifecycle hooks
---------------
Use ``on_connect`` and ``on_disconnect`` to react to connection events::

    def reset_state():
        context_store.reset_all()

    client = ConnectorClient(url="...", token="...", on_connect=reset_state)
"""

from __future__ import annotations

import json
import logging
import queue
import random
import threading
import time
from collections.abc import Callable
from typing import Any

import httpx
from httpx_sse import connect_sse

from autoplay_sdk.models import ActionsPayload, SummaryPayload

logger = logging.getLogger(__name__)

_BACKOFF_INITIAL = 1.0
_BACKOFF_MAX = 30.0
_WORKER_STOP = object()  # sentinel that tells the worker thread to exit


class ConnectorClient:
    """Subscribes to an Autoplay connector SSE stream and dispatches events to callbacks.

    The SSE receive loop and the callback execution run on separate threads.
    This means slow callbacks (DB writes, vector store upserts, LLM calls, etc.)
    never block the stream reader — the stream stays open and events keep
    arriving regardless of how long your processing takes.

    Args:
        url:             Full URL to ``GET /stream/{product_id}`` on the connector.
        token:           Unkey API key for this product (``uk_live_...``).
                         Created in the Unkey dashboard with ``external_id`` set
                         to the ``product_id``.
        max_queue_size:  Maximum number of events that can be buffered while the
                         worker thread is busy.  Events beyond this limit are
                         dropped and counted in ``dropped_count``.  Must be >= 1.
                         Default: 500.
        connect_timeout: Seconds to wait for the initial TCP connection before
                         raising an error and triggering a reconnect.  Does not
                         limit how long the stream can stay open.  Default: 10.0.
        on_connect:      Optional ``() -> None`` called each time the SSE
                         connection is successfully established (including after
                         reconnects).  Useful for resetting state.
        on_disconnect:   Optional ``() -> None`` called each time the SSE
                         connection is lost (before the reconnect backoff sleep).
    """

    def __init__(
        self,
        url: str,
        token: str = "",
        max_queue_size: int = 500,
        connect_timeout: float = 10.0,
        on_connect: Callable[[], None] | None = None,
        on_disconnect: Callable[[], None] | None = None,
    ) -> None:
        if max_queue_size < 1:
            raise ValueError(f"max_queue_size must be >= 1, got {max_queue_size!r}")
        if connect_timeout <= 0:
            raise ValueError(f"connect_timeout must be > 0, got {connect_timeout!r}")
        self._url = url
        self._token = token
        self._connect_timeout = connect_timeout
        self._on_connect_cb = on_connect
        self._on_disconnect_cb = on_disconnect
        self._actions_cb: Callable[[ActionsPayload], None] | None = None
        self._summary_cb: Callable[[SummaryPayload], None] | None = None
        self._drop_cb: Callable[[dict, int], None] | None = None
        self._running = False
        self._queue: queue.Queue = queue.Queue(maxsize=max_queue_size)
        self._dropped = 0

    # ------------------------------------------------------------------
    # Builder interface
    # ------------------------------------------------------------------

    def on_actions(self, fn: Callable[[ActionsPayload], None]) -> "ConnectorClient":
        """Register a callback for ``actions`` events.

        Called every time the connector extracts a batch of UI actions from a
        user session.  Runs on a dedicated worker thread so blocking I/O (e.g.
        vector store upserts) is safe here.

        Args:
            fn: Callable that receives a typed ``ActionsPayload``.
                Call ``payload.to_text()`` to get an embedding-ready string.

        Returns:
            self — for method chaining.
        """
        self._actions_cb = fn
        return self

    def on_summary(self, fn: Callable[[SummaryPayload], None]) -> "ConnectorClient":
        """Register a callback for ``summary`` events.

        Called when the connector's LLM summariser produces a prose summary
        of a session, replacing the raw action history.  Useful for keeping
        a compact, up-to-date context window in your RAG pipeline.

        Args:
            fn: Callable that receives a typed ``SummaryPayload``.
                Call ``payload.to_text()`` to get the prose summary string.

        Returns:
            self — for method chaining.
        """
        self._summary_cb = fn
        return self

    def on_drop(self, fn: Callable[[dict, int], None]) -> "ConnectorClient":
        """Register a callback invoked when an event is dropped due to a full queue.

        Use this to alert, log to an external system, or increment a metric.

        Args:
            fn: Called with ``(payload, total_dropped)`` where ``payload`` is
                the dropped event dict and ``total_dropped`` is the running
                drop count including this event.

        Returns:
            self — for method chaining.
        """
        self._drop_cb = fn
        return self

    # ------------------------------------------------------------------
    # Observability
    # ------------------------------------------------------------------

    @property
    def dropped_count(self) -> int:
        """Number of events dropped because the internal queue was full.

        A non-zero value means your callback is slower than the incoming event
        rate.  Consider optimising the callback or increasing ``max_queue_size``.
        """
        return self._dropped

    @property
    def queue_size(self) -> int:
        """Number of events currently waiting in the internal queue.

        Use this for periodic health logging or dashboards.  A value that grows
        steadily over time indicates the worker thread cannot keep up.
        """
        return self._queue.qsize()

    # ------------------------------------------------------------------
    # Lifecycle
    # ------------------------------------------------------------------

    def run(self) -> None:
        """Connect to the SSE stream and block until stopped or interrupted.

        Starts a background worker thread that drains the internal queue and
        calls your registered callbacks.  The SSE receive loop runs on the
        calling thread.

        Reconnects automatically with exponential backoff on any connection
        failure or network error.  Fatal HTTP errors (401, 403, 404) are
        re-raised immediately.
        """
        self._running = True
        worker = threading.Thread(
            target=self._worker_loop, daemon=True, name="autoplay-sdk-worker"
        )
        worker.start()

        backoff = _BACKOFF_INITIAL
        headers = {"Accept": "text/event-stream"}
        if self._token:
            headers["Authorization"] = f"Bearer {self._token}"
        timeout = httpx.Timeout(None, connect=self._connect_timeout)

        try:
            while self._running:
                try:
                    logger.info("autoplay_sdk: connecting to %s", self._url)
                    with httpx.Client(timeout=timeout) as client:
                        with connect_sse(
                            client, "GET", self._url, headers=headers
                        ) as source:
                            logger.info("autoplay_sdk: connected")
                            backoff = _BACKOFF_INITIAL
                            self._fire_on_connect()
                            for sse in source.iter_sse():
                                if not self._running:
                                    break
                                self._enqueue(sse)

                except httpx.HTTPStatusError as exc:
                    if exc.response.status_code in (401, 403, 404):
                        logger.error(
                            "autoplay_sdk: fatal HTTP %d — check URL and token",
                            exc.response.status_code,
                        )
                        raise
                    logger.warning(
                        "autoplay_sdk: HTTP %d — retrying in %.0fs",
                        exc.response.status_code,
                        backoff,
                    )
                    self._fire_on_disconnect()
                except KeyboardInterrupt:
                    logger.info("autoplay_sdk: interrupted")
                    break
                except Exception as exc:
                    logger.warning(
                        "autoplay_sdk: connection lost (%s) — retrying in %.0fs",
                        exc,
                        backoff,
                    )
                    self._fire_on_disconnect()

                if self._running:
                    jitter = random.uniform(0, backoff * 0.1)
                    time.sleep(backoff + jitter)
                    backoff = min(backoff * 2, _BACKOFF_MAX)
        finally:
            # Signal the worker thread to drain remaining events and exit.
            self._queue.put(_WORKER_STOP)
            worker.join(timeout=10)
            if worker.is_alive():
                logger.warning("autoplay_sdk: worker thread did not stop within 10 s")
            logger.info("autoplay_sdk: stopped")

    def run_in_background(self) -> threading.Thread:
        """Start the SSE connection in a background daemon thread.

        Returns immediately so your application can continue running.  Call
        ``stop()`` to shut down cleanly when you are done.

        Returns:
            The daemon ``Thread`` running ``run()``.  You can ``join()`` it if
            you need to wait for a clean shutdown.

        Example::

            client = ConnectorClient(url="...", token="...").on_actions(fn)
            client.run_in_background()
            # application continues; events arrive in the background
            ...
            client.stop()
        """
        t = threading.Thread(target=self.run, daemon=True, name="autoplay-sdk-main")
        t.start()
        return t

    def stop(self) -> None:
        """Signal the client to stop after the current event is processed."""
        self._running = False

    # ------------------------------------------------------------------
    # Context manager
    # ------------------------------------------------------------------

    def __enter__(self) -> "ConnectorClient":
        """Support ``with ConnectorClient(...) as client:`` usage."""
        return self

    def __exit__(self, *_: Any) -> None:
        """Call ``stop()`` automatically when exiting the ``with`` block."""
        self.stop()

    def __repr__(self) -> str:
        return (
            f"ConnectorClient(url={self._url!r}, running={self._running}, "
            f"dropped={self._dropped})"
        )

    # ------------------------------------------------------------------
    # Internal
    # ------------------------------------------------------------------

    def _fire_on_connect(self) -> None:
        if self._on_connect_cb is not None:
            try:
                self._on_connect_cb()
            except Exception as exc:
                logger.error(
                    "autoplay_sdk: on_connect callback raised: %s", exc, exc_info=True
                )

    def _fire_on_disconnect(self) -> None:
        if self._on_disconnect_cb is not None:
            try:
                self._on_disconnect_cb()
            except Exception as exc:
                logger.error(
                    "autoplay_sdk: on_disconnect callback raised: %s",
                    exc,
                    exc_info=True,
                )

    def _enqueue(self, sse: Any) -> None:
        """Parse an SSE frame and place the payload on the internal queue."""
        if sse.event == "heartbeat":
            return

        try:
            payload = json.loads(sse.data)
        except json.JSONDecodeError:
            logger.warning(
                "autoplay_sdk: could not parse event data: %s", sse.data[:200]
            )
            return

        event_type = payload.get("type")
        if event_type not in ("actions", "summary"):
            logger.debug("autoplay_sdk: unknown event type: %s", event_type)
            return

        try:
            self._queue.put_nowait(payload)
        except queue.Full:
            self._dropped += 1
            logger.warning(
                "autoplay_sdk: queue full — dropping event (type=%s, total_dropped=%d)",
                event_type,
                self._dropped,
            )
            if self._drop_cb is not None:
                try:
                    self._drop_cb(payload, self._dropped)
                except Exception as exc:
                    logger.error(
                        "autoplay_sdk: on_drop callback raised: %s", exc, exc_info=True
                    )

    def _worker_loop(self) -> None:
        """Drain the internal queue and call the registered callbacks."""
        while True:
            item = self._queue.get()
            if item is _WORKER_STOP:
                self._queue.task_done()
                break
            try:
                self._dispatch(item)
            finally:
                self._queue.task_done()

    def _dispatch(self, payload: dict) -> None:
        """Parse a raw payload dict into a typed model and call the registered callback."""
        event_type = payload.get("type")

        if event_type == "actions":
            cb = self._actions_cb
            typed: ActionsPayload | SummaryPayload = ActionsPayload.from_dict(payload)
        elif event_type == "summary":
            cb = self._summary_cb
            typed = SummaryPayload.from_dict(payload)
        else:
            return

        if cb is None:
            return

        try:
            cb(typed)
        except Exception as exc:
            logger.error(
                "autoplay_sdk: callback raised an error for event_type=%s: %s",
                event_type,
                exc,
                exc_info=True,
            )
