"""Auto-injected capability that drains the pending message queue at appropriate times."""

from __future__ import annotations

from typing import TYPE_CHECKING, Any

from pydantic_ai._agent_graph import ModelRequestNode
from pydantic_ai._enqueue import PendingMessage, PendingMessageQueue
from pydantic_ai._utils import fill_run_metadata
from pydantic_ai.capabilities.abstract import AbstractCapability, CapabilityOrdering
from pydantic_ai.exceptions import UserError
from pydantic_ai.messages import EnqueuedMessagesEvent, ModelMessage, ModelRequest
from pydantic_ai.tools import RunContext
from pydantic_graph import End

if TYPE_CHECKING:
    from pydantic_ai import _agent_graph
    from pydantic_ai.models import ModelRequestContext
    from pydantic_ai.result import FinalResult


def _queue(ctx: RunContext[Any]) -> PendingMessageQueue:
    # Graph runs always carry the thread-safe queue. Other `RunContext` producers (durable-exec
    # guards, synthetic contexts) never reach this drain.
    queue = ctx.pending_messages
    assert isinstance(queue, PendingMessageQueue)
    return queue


def _stamped_messages(
    pending: PendingMessage,
    *,
    fallback_run_id: str | None,
    fallback_conversation_id: str | None,
) -> list[ModelMessage]:
    """Stamp a pending message's messages' `timestamp` / `run_id` / `conversation_id` where unset.

    Each [`PendingMessage`][pydantic_ai._enqueue.PendingMessage] carries one or more built
    [`ModelMessage`][pydantic_ai.messages.ModelMessage]s (assembled at enqueue time by
    [`PendingMessage.from_content`][pydantic_ai._enqueue.PendingMessage.from_content]); this only
    fills in framework-tracked metadata that the producer left unset, so producer-supplied values
    are preserved.
    """
    messages: list[ModelMessage] = []
    for message in pending.messages:
        fill_run_metadata(message, run_id=fallback_run_id, conversation_id=fallback_conversation_id)
        messages.append(message)
    return messages


class PendingMessageDrainCapability(AbstractCapability[Any]):
    """Drains the pending message queue at appropriate times.

    - `'asap'` messages drain at the earliest opportunity: into the next
      [`ModelRequest`][pydantic_ai.messages.ModelRequest] via `before_model_request`,
      or — if the agent would otherwise terminate — redirected through a new
      `ModelRequestNode` at the end of the run.
    - `'when_idle'` messages drain only when the agent would otherwise terminate
      and no `'asap'` messages remain, after any `'asap'` redirect.

    This capability is always auto-injected and placed outermost via
    [`CapabilityOrdering`][pydantic_ai.capabilities.abstract.CapabilityOrdering].
    Final draining happens at the graph-advancement seam after the full capability
    hook chain, so user ordering constraints cannot make the queue close before a
    hook has had a chance to redirect an apparent [`End`][pydantic_graph.End].
    """

    def get_ordering(self) -> CapabilityOrdering:
        return CapabilityOrdering(position='outermost')

    @classmethod
    def get_serialization_name(cls) -> str | None:
        return None  # not spec-constructible (internal, auto-injected)

    async def before_model_request(
        self,
        ctx: RunContext[Any],
        request_context: ModelRequestContext,
    ) -> ModelRequestContext:
        """Drain `'asap'` messages into the upcoming model request.

        Each drained request is appended to both `request_context.messages` (so the model
        sees it this step) and `ctx.messages` (so it persists in the agent's message
        history). Stamps `timestamp`/`run_id`/`conversation_id` if the producer didn't —
        `ModelRequestNode.run()` only stamps `self.request` (the current node's request),
        and capabilities downstream of us might append more messages, so we can't rely on
        that fixup.

        Emits one [`EnqueuedMessagesEvent`][pydantic_ai.messages.EnqueuedMessagesEvent] per drained
        [`enqueue`][pydantic_ai.tools.RunContext.enqueue] call, in enqueue order, describing the
        messages exactly as delivered here.
        """
        drained = _queue(ctx).pop_priority('asap')
        for pending in drained:
            messages = _stamped_messages(
                pending, fallback_run_id=ctx.run_id, fallback_conversation_id=ctx.conversation_id
            )
            request_context.messages.extend(messages)
            ctx.messages.extend(messages)
            ctx._emit_event(EnqueuedMessagesEvent(enqueue_id=pending.enqueue_id, messages=tuple(messages)))  # pyright: ignore[reportPrivateUsage]
        return request_context


def drain_pending_messages_at_end(
    ctx: RunContext[Any],
    result: _agent_graph.AgentNode[Any, Any] | End[FinalResult[Any]],
) -> _agent_graph.AgentNode[Any, Any] | End[FinalResult[Any]]:
    """Drain pending messages after all capability hooks if the agent would terminate.

    Drain `'asap'` messages first (anything that arrived after the most recent
    `before_model_request` and would otherwise be lost), then `'when_idle'` messages.
    Each priority is appended independently so the history keeps the priority split
    visible (matches pi-mono's separate steering / follow-up turns). On the wire,
    `_clean_message_history` re-merges adjacent requests with compatible instructions,
    so the model still sees one turn.

    The last resulting request becomes the redirect
    [`ModelRequestNode`][pydantic_ai._agent_graph.ModelRequestNode]'s request; any
    earlier ones are appended to `ctx.messages` so they appear in history before the
    redirect. Emits one [`EnqueuedMessagesEvent`][pydantic_ai.messages.EnqueuedMessagesEvent]
    per drained [`enqueue`][pydantic_ai.tools.RunContext.enqueue] call, in enqueue order.
    """
    if not isinstance(result, End):
        return result

    # Pi-mono parity: drain `'asap'` first so anything that arrived during the
    # final step (e.g. a background task completing while the model produced
    # its final response) gets delivered before `'when_idle'` messages, and the
    # agent gets another turn rather than terminating with the message lost.
    leftover_asap, when_idle = _queue(ctx).drain_at_end()
    if not leftover_asap and not when_idle:
        return result

    drained = [*leftover_asap, *when_idle]
    stamped = [
        (
            pending,
            _stamped_messages(pending, fallback_run_id=ctx.run_id, fallback_conversation_id=ctx.conversation_id),
        )
        for pending in drained
    ]
    messages = [message for _, pending_messages in stamped for message in pending_messages]
    # `final` becomes the redirect node's request; `ModelRequestNode._prepare_request`
    # will re-stamp it during the graph lifecycle. `_stamped_messages` already
    # stamped it, which is harmless (the lifecycle stamp overwrites). `from_content`
    # guarantees each `PendingMessage` ends in a `ModelRequest`, but a producer can
    # construct `PendingMessage` (or mutate `RunContext.pending_messages`) directly, so
    # we check rather than assert. Every message except `final` is appended to history
    # before the redirect.
    final = messages[-1]
    if not isinstance(final, ModelRequest):
        raise UserError(
            'Enqueued content must end with a `ModelRequest` so the agent has a request to respond to, '
            f'but the last queued message is a `{type(final).__name__}`.'
        )
    for pending, pending_messages in stamped:
        for message in pending_messages:
            if message is not final:
                ctx.messages.append(message)
        ctx._emit_event(  # pyright: ignore[reportPrivateUsage]
            EnqueuedMessagesEvent(enqueue_id=pending.enqueue_id, messages=tuple(pending_messages))
        )
    return ModelRequestNode(request=final)