"""autoplay_sdk.async_client — Async-native SSE client for Autoplay connectors.

Identical feature set to ``ConnectorClient`` but built for ``asyncio``-first
environments: LangChain, LlamaIndex, FastAPI backends, and any RAG pipeline
that uses ``async/await``.

Callbacks are coroutines (``async def``).  This means you can ``await``
vector store upserts, LLM embedding calls, or database writes directly inside
your callback without blocking the event loop or spawning extra threads.

Quickstart::

    import asyncio
    from autoplay_sdk import AsyncConnectorClient, ActionsPayload

    async def on_actions(payload: ActionsPayload):
        text = payload.to_text()
        # await your vector store, LLM, or database here
        print(text)

    async def main():
        async with AsyncConnectorClient(url="https://host/stream/my_product",
                                        token="uk_live_...") as client:
            client.on_actions(on_actions)
            await client.run()

    asyncio.run(main())

Background usage (non-blocking inside an existing event loop)::

    client = AsyncConnectorClient(url="...", token="...")
    client.on_actions(on_actions)
    task = client.run_in_background()  # returns asyncio.Task

    # do other async work ...
    await asyncio.sleep(60)

    client.stop()
    await task

Reconnection
------------
Automatically reconnects with exponential backoff (1 s → 30 s) on any
network error.  Fatal HTTP errors (401, 403, 404) are re-raised immediately.

Per-session task isolation
--------------------------
Each incoming payload is dispatched as an independent ``asyncio.Task`` keyed
by ``session_id``, mirroring the event connector's ``ProductWorkerThread``
architecture.  This means:

- A slow callback for session A never delays delivery for session B.
- Each session has its own ``asyncio.Semaphore`` (``session_concurrency``)
  limiting the number of concurrent in-flight dispatches for that session,
  providing backpressure without cross-session interference.

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

    async def reset_state():
        await context_store.reset_all()

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

from __future__ import annotations

import asyncio
import json
import logging
import random
from collections.abc import Callable
from typing import Any

import httpx
from httpx_sse import aconnect_sse

from autoplay_sdk.models import ActionsPayload, SummaryPayload

logger = logging.getLogger(__name__)

_BACKOFF_INITIAL = 1.0
_BACKOFF_MAX = 30.0

# Callbacks may be sync or async; coroutine detection happens at dispatch time.
ActionCallback = Callable[[ActionsPayload], Any]
SummaryCallback = Callable[[SummaryPayload], Any]
LifecycleCallback = Callable[[], Any]


class AsyncConnectorClient:
    """Async SSE client that dispatches typed payload models to async callbacks.

    Designed for ``asyncio``-first RAG pipelines where callbacks need to
    ``await`` embedding APIs, vector stores, or databases without blocking.

    Each payload is dispatched as an independent ``asyncio.Task`` keyed by
    ``session_id`` so sessions never block each other.  A per-session
    ``asyncio.Semaphore`` (``session_concurrency``) limits concurrent in-flight
    dispatches per session.

    Args:
        url:                 Full URL to ``GET /stream/{product_id}`` on the connector.
        token:               Unkey API key for this product (``uk_live_...``).
        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.
        session_concurrency: Maximum number of concurrent in-flight callback tasks
                             per session.  A session that exceeds this will wait
                             for an earlier dispatch to finish before processing
                             the next payload.  Default: 4.
        on_connect:          Optional async or sync ``() -> None`` called each time
                             the SSE connection is successfully established (including
                             after reconnects).  Useful for resetting state.
        on_disconnect:       Optional async or sync ``() -> None`` called each time
                             the SSE connection is lost (before the reconnect backoff
                             sleep).
    """

    def __init__(
        self,
        url: str,
        token: str = "",
        connect_timeout: float = 10.0,
        session_concurrency: int = 4,
        on_connect: LifecycleCallback | None = None,
        on_disconnect: LifecycleCallback | None = None,
    ) -> None:
        if connect_timeout <= 0:
            raise ValueError(f"connect_timeout must be > 0, got {connect_timeout!r}")
        if session_concurrency < 1:
            raise ValueError(
                f"session_concurrency must be >= 1, got {session_concurrency!r}"
            )
        self._url = url
        self._token = token
        self._connect_timeout = connect_timeout
        self._session_concurrency = session_concurrency
        self._on_connect_cb = on_connect
        self._on_disconnect_cb = on_disconnect
        self._actions_cb: ActionCallback | None = None
        self._summary_cb: SummaryCallback | None = None
        self._running = False
        # Per-session semaphores — created lazily, keyed by session_id.
        # Limits concurrent in-flight callback tasks per session independently.
        self._session_semaphores: dict[str, asyncio.Semaphore] = {}

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

    def on_actions(self, fn: ActionCallback) -> "AsyncConnectorClient":
        """Register an async callback for ``actions`` events.

        Args:
            fn: ``async def fn(payload: ActionsPayload) -> None``
                Call ``payload.to_text()`` for an embedding-ready string.

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

    def on_summary(self, fn: SummaryCallback) -> "AsyncConnectorClient":
        """Register an async callback for ``summary`` events.

        Args:
            fn: ``async def fn(payload: SummaryPayload) -> None``
                Call ``payload.to_text()`` for the prose summary string.

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

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

    async def run(self) -> None:
        """Connect to the SSE stream and process events until stopped.

        Reconnects automatically with exponential backoff on network failures.
        Fatal HTTP errors (401, 403, 404) are raised immediately.

        Intended to be the main coroutine in your async application::

            await client.run()
        """
        self._running = True
        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)

        while self._running:
            try:
                logger.info("autoplay_sdk: connecting to %s", self._url)
                async with httpx.AsyncClient(timeout=timeout) as http:
                    async with aconnect_sse(
                        http, "GET", self._url, headers=headers
                    ) as source:
                        logger.info("autoplay_sdk: connected")
                        backoff = _BACKOFF_INITIAL
                        await self._fire_on_connect()
                        async for sse in source.aiter_sse():
                            if not self._running:
                                return
                            await self._handle(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,
                )
                await self._fire_on_disconnect()
            except asyncio.CancelledError:
                logger.info("autoplay_sdk: cancelled")
                self._running = False
                return
            except Exception as exc:
                logger.warning(
                    "autoplay_sdk: connection lost (%s) — retrying in %.0fs",
                    exc,
                    backoff,
                )
                await self._fire_on_disconnect()

            if self._running:
                jitter = random.uniform(0, backoff * 0.1)
                await asyncio.sleep(backoff + jitter)
                backoff = min(backoff * 2, _BACKOFF_MAX)

        logger.info("autoplay_sdk: stopped")

    def run_in_background(self) -> asyncio.Task:
        """Schedule ``run()`` as a background ``asyncio.Task``.

        Returns immediately.  The task runs concurrently with other
        coroutines in the current event loop.  Call ``stop()`` and then
        ``await`` the returned task to shut down cleanly.

        Returns:
            ``asyncio.Task`` wrapping ``run()``.

        Example::

            task = client.run_in_background()
            await do_other_async_work()
            client.stop()
            await task
        """
        return asyncio.ensure_future(self.run())

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

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

    async def __aenter__(self) -> "AsyncConnectorClient":
        """Support ``async with AsyncConnectorClient(...) as client:`` usage."""
        return self

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

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

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

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

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

    async def _handle(self, sse: Any) -> None:
        """Parse one SSE frame and dispatch it to the registered callback."""
        if sse.event == "heartbeat":
            return

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

        event_type = raw.get("type")

        if event_type == "actions":
            typed: ActionsPayload | SummaryPayload = ActionsPayload.from_dict(raw)
        elif event_type == "summary":
            typed = SummaryPayload.from_dict(raw)
        else:
            logger.debug("autoplay_sdk: unknown event type: %s", event_type)
            return

        session_id = getattr(typed, "session_id", None) or "_nosession"
        await self._dispatch(session_id, typed)

    async def _dispatch(
        self, session_id: str, typed: ActionsPayload | SummaryPayload
    ) -> None:
        """Invoke the registered callback under the per-session semaphore.

        Each session has its own ``asyncio.Semaphore`` (``session_concurrency``)
        so slow callbacks for one session never delay processing for another.
        """
        sem = self._session_semaphores.get(session_id)
        if sem is None:
            sem = asyncio.Semaphore(self._session_concurrency)
            self._session_semaphores[session_id] = sem

        cb: ActionCallback | SummaryCallback | None
        if isinstance(typed, ActionsPayload):
            cb = self._actions_cb
            event_type = "actions"
        else:
            cb = self._summary_cb
            event_type = "summary"

        if cb is None:
            return

        async with sem:
            try:
                result = cb(typed)
                if asyncio.iscoroutine(result):
                    await result
            except Exception as exc:
                logger.error(
                    "autoplay_sdk: callback raised an error for event_type=%s session=%s: %s",
                    event_type,
                    session_id,
                    exc,
                    exc_info=True,
                )
