from __future__ import annotations

import asyncio
from collections.abc import AsyncIterator, Awaitable, Callable, Iterator, Mapping
from types import MappingProxyType, TracebackType
from typing import TYPE_CHECKING, Any

from langchain_core._api import beta

from langgraph.stream._convert import convert_to_protocol_event
from langgraph.stream._mux import StreamMux
from langgraph.stream._types import ProtocolEvent

if TYPE_CHECKING:
    from langchain_core.language_models.chat_model_stream import (
        AsyncChatModelStream,
        ChatModelStream,
    )

    from langgraph.stream.stream_channel import StreamChannel
    from langgraph.stream.transformers import LifecyclePayload, SubgraphStatus


def _drive_until_done(pump: Callable[[], bool]) -> None:
    """Call the sync pump until it returns False."""
    while pump():
        pass


async def _adrive_until_done(pump: Callable[[], Awaitable[bool]]) -> None:
    """Call the async pump until it returns False."""
    while await pump():
        pass


@beta(message="The v3 streaming protocol on Pregel is experimental.")
class GraphRunStream:
    """Sync run stream with caller-driven pumping.

    The caller's iteration on any projection (`values`, `messages`,
    raw events, or `output`) drives the graph forward. No background
    thread is used — the caller's `for` loop is the pump.

    Projections are single-consumer — iterating `run.values` twice
    raises. Use `projection.tee(n)` if you genuinely need fan-out.

    All transformer projections live in `extensions`. Native transformer
    projections (those with `_native = True`) are also set as direct
    attributes on this instance (e.g. `run.values`, `run.messages`).

    !!! warning

        Returned by `Pregel.stream_events(version="v3")`, which is
        experimental and may change.
    """

    # Native projections always registered by `stream_events(version="v3")`.
    # Attached dynamically by the `setattr` loop in `__init__`; declared here
    # so type checkers see them. Opt-in native projections (`updates`,
    # `custom`, `checkpoints`, `debug`, `tasks`) are only present when their
    # transformer is registered, so they are reached via `extensions[...]`.
    values: StreamChannel[dict[str, Any]]
    messages: StreamChannel[ChatModelStream]
    lifecycle: StreamChannel[LifecyclePayload]
    subgraphs: StreamChannel[SubgraphRunStream]

    def __init__(
        self,
        graph_iter: Iterator[Any] | None,
        mux: StreamMux,
        *,
        wire_pump: bool = True,
    ) -> None:
        """Initialize the run stream.

        Args:
            graph_iter: Pull-based iterator over the graph's stream,
                or `None` for nested run streams whose pump is driven
                by an outer run (e.g. `SubgraphRunStream`).
            mux: The StreamMux owning projections and the main log.
            wire_pump: When True (default), bind `_pump_next` as the
                mux's pump callable. Subclasses that inherit a parent
                pump via `StreamMux._make_child` should pass False to
                preserve the parent binding.
        """
        self._graph_iter = graph_iter
        self._mux = mux
        self.extensions: Mapping[str, Any] = MappingProxyType(mux.extensions)
        self._exhausted = False
        self._latest: dict[str, Any] | None = None
        self._interrupted = False
        self._interrupts: list[Any] = []
        self._scope_list: list[str] = list(mux.scope)
        for key in mux.native_keys:
            setattr(self, key, mux.extensions[key])
        if wire_pump:
            self._wire_request_more(mux)

    def _wire_request_more(self, mux: StreamMux) -> None:
        """Wire the sync pull callback through the mux.

        Routing through `mux.bind_pump` (rather than walking
        projections directly here) lets child mini-muxes built by
        `mux._make_child(...)` inherit the same pump callable, so
        cursors on a subgraph handle's projections drive the root
        pump just like cursors on `run.values` do.
        """
        mux.bind_pump(self._pump_next)

    def _observe_event(self, event: ProtocolEvent) -> None:
        """Track values-event state for output/interrupted/interrupts."""
        if event["method"] != "values":
            return
        params = event["params"]
        if params["namespace"] != self._scope_list:
            return
        self._latest = params["data"]
        interrupts = params.get("interrupts", ())
        if interrupts:
            self._interrupted = True
            self._interrupts.extend(interrupts)

    def _pump_next(self) -> bool:
        """Pull one event from the graph and push it through the mux.

        Returns:
            True if an event was pulled, False if the graph is exhausted
            or has raised. Always False when constructed with
            `graph_iter=None` (the run is driven by an outer pump).
        """
        if self._exhausted or self._graph_iter is None:
            return False
        try:
            part = next(self._graph_iter)
            event = convert_to_protocol_event(part)
            self._observe_event(event)
            self._mux.push(event)
            return True
        except StopIteration:
            self._mux.close()
            self._exhausted = True
            return False
        except Exception as e:
            self._mux.fail(e)
            self._exhausted = True
            return False

    def abort(self) -> None:
        """Stop the run early.

        Closes the underlying graph iterator (propagating `GeneratorExit`
        so in-flight nodes and subgraphs are cancelled), closes the mux,
        and marks the stream exhausted. Idempotent.
        """
        if self._exhausted:
            return
        self._exhausted = True
        graph_iter = self._graph_iter
        self._graph_iter = None
        if (
            graph_iter is not None
            and (close := getattr(graph_iter, "close", None)) is not None
        ):
            try:
                close()
            except Exception:
                pass
        try:
            self._mux.close()
        except Exception:
            pass

    def __enter__(self) -> GraphRunStream:
        return self

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc: BaseException | None,
        tb: TracebackType | None,
    ) -> None:
        self.abort()

    @property
    def output(self) -> dict[str, Any] | None:
        """Drive the run to completion and return the final state."""
        _drive_until_done(self._pump_next)
        if (err := self._mux._events._error) is not None:
            raise err
        return self._latest

    @property
    def interrupted(self) -> bool:
        """Drive the run to completion, then return whether it was
        interrupted.

        Raises:
            BaseException: If the run ended with an error.
        """
        _drive_until_done(self._pump_next)
        if (err := self._mux._events._error) is not None:
            raise err
        return self._interrupted

    @property
    def interrupts(self) -> list[Any]:
        """Drive the run to completion, then return interrupt payloads.

        Raises:
            BaseException: If the run ended with an error.
        """
        _drive_until_done(self._pump_next)
        if (err := self._mux._events._error) is not None:
            raise err
        return self._interrupts

    def __iter__(self) -> Iterator[ProtocolEvent]:
        """Subscribe to the main event log and iterate protocol events."""
        return iter(self._mux._events)

    def interleave(self, *names: str) -> Iterator[tuple[str, Any]]:
        """Iterate multiple projections in arrival order, yielding ``(name, item)``.

        Items are ordered by a monotonic push stamp assigned when each
        transformer pushes into its `StreamChannel`. This gives strict
        arrival ordering across projections, unlike round-robin.

        Args:
            *names: Projection keys to interleave. Must match keys in
                ``extensions``.

        Yields:
            ``(name, item)`` tuples in arrival order across the named
            projections.

        Each named channel is locked for the duration of iteration and
        released when the generator completes, is closed, or raises.
        Channels cannot be subscribed concurrently — use `.tee(n)` if
        you need fan-out.

        Raises:
            KeyError: If a name doesn't match a registered projection.

        Example:
            ```python
            for name, item in run.interleave("messages", "values"):
                if name == "messages":
                    print("msg:", item)
                else:
                    print("val:", item)
            ```
        """
        from langgraph.stream.stream_channel import StreamChannel

        channels: dict[str, StreamChannel[Any]] = {}
        try:
            for name in names:
                ch = self.extensions[name]
                if not isinstance(ch, StreamChannel):
                    raise TypeError(
                        f"interleave() requires StreamChannel projections, "
                        f"got {type(ch).__name__} for {name!r}"
                    )
                if ch._is_async is None:
                    raise TypeError(
                        f"StreamChannel {name!r} has not been bound yet. "
                        "Register the transformer with a StreamMux first."
                    )
                if ch._is_async:
                    raise TypeError(
                        f"StreamChannel {name!r} is bound to async mode — "
                        "sync interleave() cannot consume async channels."
                    )
                if ch._subscribed:
                    raise RuntimeError(
                        f"StreamChannel {name!r} already has a subscriber; "
                        "use .tee(n) for fan-out."
                    )
                ch._subscribed = True
                channels[name] = ch

            done: set[str] = set()

            while len(done) < len(channels):
                best: tuple[int, str] | None = None
                for name, ch in channels.items():
                    if name in done:
                        continue
                    if ch._closed and not ch._items:
                        if ch._error is not None:
                            raise ch._error
                        done.add(name)
                        continue
                    if ch._items:
                        stamp = ch._items[0][0]
                        if best is None or stamp < best[0]:
                            best = (stamp, name)

                if best is not None:
                    _stamp, item = channels[best[1]]._items.popleft()
                    yield (best[1], item)
                else:
                    pump = self._mux._pump_fn
                    if pump is None or not pump():
                        before = len(done)
                        for name, ch in channels.items():
                            if name not in done and not ch._items:
                                if ch._closed:
                                    if ch._error is not None:
                                        raise ch._error
                                    done.add(name)
                        if len(done) == before:
                            break
        finally:
            for ch in channels.values():
                ch._subscribed = False


@beta(message="The v3 streaming protocol on Pregel is experimental.")
class AsyncGraphRunStream:
    """Async run stream with caller-driven pumping.

    Async iteration on any projection drives the graph forward — there
    is no background task. Concurrent consumers share a single-flight
    pump via an `asyncio.Lock`, so each awaiting cursor contributes one
    event per acquisition. Backpressure comes from the logs: when a
    subscribed log's buffer reaches `maxlen`, `apush` awaits the
    subscriber to drain, which holds back the pump and paces the graph.

    Projections are single-consumer — a second `aiter(run.values)`
    raises. Use `projection.tee(n)` for fan-out.

    Use as an async context manager to guarantee clean shutdown on
    early exit:

    ```python
    async with await handler.astream(input) as run:
        async for msg in run.messages:
            ...
    ```

    !!! warning

        Awaited from `Pregel.astream_events(version="v3")`, which is
        experimental and may change.
    """

    # Native projections always registered by `astream_events(version="v3")`.
    # Attached dynamically by the `setattr` loop in `__init__`; declared here
    # so type checkers see them. Opt-in native projections (`updates`,
    # `custom`, `checkpoints`, `debug`, `tasks`) are only present when their
    # transformer is registered, so they are reached via `extensions[...]`.
    values: StreamChannel[dict[str, Any]]
    messages: StreamChannel[AsyncChatModelStream]
    lifecycle: StreamChannel[LifecyclePayload]
    subgraphs: StreamChannel[AsyncSubgraphRunStream]

    def __init__(
        self,
        graph_aiter: AsyncIterator[Any] | None,
        mux: StreamMux,
        *,
        wire_pump: bool = True,
    ) -> None:
        """Initialize the async run stream.

        Args:
            graph_aiter: Async iterator over the graph's stream, or
                `None` for nested run streams whose pump is driven by
                an outer run (e.g. `AsyncSubgraphRunStream`).
            mux: The StreamMux owning projections and the main log.
            wire_pump: When True (default), bind `_apump_next` as the
                mux's async pump callable. Subclasses that inherit a
                parent pump via `StreamMux._make_child` should pass
                False to preserve the parent binding.
        """
        self._graph_aiter = graph_aiter
        self._mux = mux
        self.extensions: Mapping[str, Any] = MappingProxyType(mux.extensions)
        self._exhausted = False
        self._latest: dict[str, Any] | None = None
        self._interrupted = False
        self._interrupts: list[Any] = []
        self._scope_list: list[str] = list(mux.scope)
        self._pump_cond = asyncio.Condition()
        self._pumping = False
        self._anext_task: asyncio.Future[Any] | None = None
        self._aborting = False
        for key in mux.native_keys:
            setattr(self, key, mux.extensions[key])
        if wire_pump:
            self._wire_arequest_more(mux)

    def _observe_event(self, event: ProtocolEvent) -> None:
        """Track values-event state for output/interrupted/interrupts."""
        if event["method"] != "values":
            return
        params = event["params"]
        if params["namespace"] != self._scope_list:
            return
        self._latest = params["data"]
        interrupts = params.get("interrupts", ())
        if interrupts:
            self._interrupted = True
            self._interrupts.extend(interrupts)

    def _wire_arequest_more(self, mux: StreamMux) -> None:
        """Wire the async pull callback through the mux.

        Mirrors `_wire_request_more`: routing through
        `mux.bind_apump` lets child mini-muxes inherit the pump
        callable so cursors on subgraph handles drive the root
        pump.
        """
        mux.bind_apump(self._apump_next)

    async def _apump_next(self) -> bool:
        """Drive one pump step, or wait for the active pumper to drive one.

        "Take-a-number" semantics: at most one task at a time calls
        `graph_aiter.__anext__()` (asyncio iterators can't be advanced
        concurrently). Other callers wait on a Condition that the
        active pumper notifies after each step. This lets a "passive"
        consumer — one whose projection's buffer is being filled by the
        active pumper's push — wake up as soon as its data lands,
        instead of queueing on the pump and only observing its data one
        graph event late.

        `except Exception` is intentional — `CancelledError` and other
        `BaseException` subclasses propagate, matching asyncio's
        cancellation contract.

        Returns:
            True if a pump step completed (by this task or another),
            False if the graph is exhausted.
        """
        async with self._pump_cond:
            if self._exhausted or self._graph_aiter is None:
                return False
            if self._pumping:
                # Another task is pumping; wait for its progress signal.
                await self._pump_cond.wait()
                return not self._exhausted
            self._pumping = True

        try:
            try:
                # Run the pull as a child task so `abort()` can cancel it
                # mid-flight. Cancelling propagates `CancelledError` into the
                # graph generator frame -> Pregel loop -> nested subgraph
                # nodes, which a bare `aclose()` cannot do while the generator
                # is running ("asynchronous generator is already running").
                self._anext_task = asyncio.ensure_future(self._graph_aiter.__anext__())
                try:
                    part = await self._anext_task
                except asyncio.CancelledError:
                    if self._aborting:
                        # Abort-initiated cancel: stop gracefully.
                        self._exhausted = True
                        return False
                    # Genuine external cancel of this task: also stop the
                    # in-flight pull, then propagate.
                    self._anext_task.cancel()
                    raise
                finally:
                    self._anext_task = None
                event = convert_to_protocol_event(part)
                self._observe_event(event)
                await self._mux.apush(event)
                return True
            except StopAsyncIteration:
                self._exhausted = True
                await self._mux.aclose()
                return False
            except Exception as e:
                self._exhausted = True
                await self._mux.afail(e)
                return False
        finally:
            async with self._pump_cond:
                self._pumping = False
                self._pump_cond.notify_all()

    async def abort(self) -> None:
        """Stop the run early.

        Marks the stream exhausted and wakes any pump-waiters. Cancels an
        in-flight pull if one is running, then closes the underlying graph
        iterator, so running nodes and nested subgraphs are cancelled
        whether or not a pump is mid-pull. Closes the mux; any `apush`
        blocked on backpressure wakes and returns without appending.
        Idempotent.
        """
        async with self._pump_cond:
            if self._exhausted:
                return
            self._exhausted = True
            self._aborting = True
            graph_aiter = self._graph_aiter
            self._graph_aiter = None
            anext_task = self._anext_task
            self._pump_cond.notify_all()
        # If a pump is mid-pull, cancel it so the cancellation propagates
        # into running nodes and nested subgraphs. Once it settles the
        # generator is no longer running, so the `aclose()` below is a safe
        # final cleanup (and handles the no-in-flight-pull case directly).
        if anext_task is not None and not anext_task.done():
            anext_task.cancel()
            try:
                await anext_task
            except (asyncio.CancelledError, Exception):
                pass
        if (
            graph_aiter is not None
            and (aclose := getattr(graph_aiter, "aclose", None)) is not None
        ):
            try:
                await aclose()
            except Exception:
                pass
        try:
            await self._mux.aclose()
        except Exception:
            pass

    async def __aenter__(self) -> AsyncGraphRunStream:
        return self

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc: BaseException | None,
        tb: TracebackType | None,
    ) -> None:
        await self.abort()

    async def output(self) -> dict[str, Any] | None:
        """Drive the run to completion and return the final state.

        Methods (not properties) on the async lane so `run.output`
        without `await` raises at type-check time instead of silently
        yielding a coroutine object.

        Example:
            ```python
            output = await run.output()
            ```

        Raises:
            BaseException: If the run ended with an error.
        """
        await _adrive_until_done(self._apump_next)
        if (err := self._mux._events._error) is not None:
            raise err
        return self._latest

    async def interrupted(self) -> bool:
        """Drive the run to completion and return whether it was
        interrupted.

        Raises:
            BaseException: If the run ended with an error.
        """
        await _adrive_until_done(self._apump_next)
        if (err := self._mux._events._error) is not None:
            raise err
        return self._interrupted

    async def interrupts(self) -> list[Any]:
        """Drive the run to completion and return interrupt payloads.

        Raises:
            BaseException: If the run ended with an error.
        """
        await _adrive_until_done(self._apump_next)
        if (err := self._mux._events._error) is not None:
            raise err
        return self._interrupts

    def __aiter__(self) -> AsyncIterator[ProtocolEvent]:
        """Subscribe to the main event log and iterate protocol events."""
        return self._mux._events.__aiter__()


class _SubgraphRunStreamMixin:
    """Subgraph metadata + parent-pump delegation shared by both lanes.

    Inherits from `GraphRunStream` (or `AsyncGraphRunStream`) with
    `graph_iter=None` + `wire_pump=False` — the mini-mux is driven
    by the parent's pump (inherited via `StreamMux._make_child`), and
    the handle never pulls upstream itself. Pump-driving methods
    delegate to the parent pump so `handle.output` and friends drive
    the root run.

    Subclasses set the parent pump function captured at construction
    (`_parent_pump_fn` / `_parent_apump_fn`) and override
    `_pump_next` / `_apump_next` to delegate to it.

    Status is updated in place by `SubgraphTransformer`. Iterate
    `run.subgraphs` to receive handles as subgraphs spawn, then
    drill into projections inside the loop body **before** the next
    pump cycle — same lazy-subscribe constraint as root projections.
    """

    path: tuple[str, ...]
    graph_name: str | None
    trigger_call_id: str | None
    status: SubgraphStatus
    error: str | None
    _seen_terminal: bool


class SubgraphRunStream(GraphRunStream, _SubgraphRunStreamMixin):
    """Sync handle for a discovered subgraph (extends `GraphRunStream`)."""

    def __init__(
        self,
        mux: StreamMux,
        *,
        path: tuple[str, ...],
        graph_name: str | None = None,
        trigger_call_id: str | None = None,
    ) -> None:
        # Capture the parent-inherited pump before super().__init__
        # touches anything; we delegate to it from `_pump_next`.
        self._parent_pump_fn: Callable[[], bool] | None = mux._pump_fn
        super().__init__(
            graph_iter=None,
            mux=mux,
            wire_pump=False,
        )
        self.path = path
        self.graph_name = graph_name
        self.trigger_call_id = trigger_call_id
        self.status = "started"
        self.error = None
        self._seen_terminal = False

    def _pump_next(self) -> bool:
        """Delegate to the parent's pump.

        Cursors on this handle's projections call here when their
        buffers empty. Driving the parent fans events into our
        mini-mux, transparently advancing the whole run.
        """
        if (
            self._exhausted
            or self._seen_terminal
            or self._mux._events._closed
            or self._parent_pump_fn is None
        ):
            return False
        return self._parent_pump_fn()


class AsyncSubgraphRunStream(AsyncGraphRunStream, _SubgraphRunStreamMixin):
    """Async handle for a discovered subgraph (extends `AsyncGraphRunStream`)."""

    def __init__(
        self,
        mux: StreamMux,
        *,
        path: tuple[str, ...],
        graph_name: str | None = None,
        trigger_call_id: str | None = None,
    ) -> None:
        self._parent_apump_fn: Callable[[], Awaitable[bool]] | None = mux._apump_fn
        super().__init__(
            graph_aiter=None,
            mux=mux,
            wire_pump=False,
        )
        self.path = path
        self.graph_name = graph_name
        self.trigger_call_id = trigger_call_id
        self.status = "started"
        self.error = None
        self._seen_terminal = False

    async def _apump_next(self) -> bool:
        """Delegate to the parent's async pump."""
        if (
            self._exhausted
            or self._seen_terminal
            or self._mux._events._closed
            or self._parent_apump_fn is None
        ):
            return False
        return await self._parent_apump_fn()
