"""PydanticAI bridge — Agent mode and pydantic_graph Graph mode.

One bridge class, two input modes.  Duck-types at construction and
routes to the appropriate internal path.
"""

from __future__ import annotations

import asyncio
import inspect
import json
import logging
import shlex
import time
from collections.abc import AsyncIterator, Callable, Mapping, Sequence
from contextlib import ExitStack
from dataclasses import replace
from typing import Any, ClassVar
from urllib.parse import unquote
from uuid import uuid4

from easycat.cancel import CancelToken
from easycat.integrations.agents._context import normalize_context_messages
from easycat.integrations.agents._helpers import (
    aclose_quietly,
    bridge_version_info,
    record_usage_from_result,
    resolve_model_name,
    split_replacement_by_original_parts,
)
from easycat.integrations.agents._pydantic_ai_events import translate_event
from easycat.integrations.agents._state_serialization import serialize_framework_state
from easycat.integrations.agents._text_stream import AgentTextStream, is_text_event
from easycat.integrations.agents.base import (
    AgentBridgeEvent,
    AgentRecorder,
    AgentTurnInput,
    BridgeConfigurationError,
    BridgeInputError,
    CancellationMode,
    CommitRule,
    ExecutionCursor,
    FrameworkStateSnapshot,
    InterruptionPlan,
    UnitKind,
    run_interruption_journal_protocol,
)
from easycat.runtime.records import ErrorInfo

logger = logging.getLogger(__name__)


def _completion_text(accumulated: str, structured_output: Any) -> str:
    """Prefer streamed text, falling back to a speakable structured result."""
    if accumulated:
        return accumulated
    if structured_output is None:
        return ""
    return str(structured_output)


class _RunPartIndexAllocator:
    """Remap stream-local text part indexes onto one run-global namespace.

    PydanticAI part indexes are local to a single streamed model response.
    A run that spans several responses (tool-call rounds in ``agent.iter()``,
    per-agent streams inside a graph) restarts at index 0 for each response,
    so forwarding raw indexes would misreport a later response's text as a
    replacement of an earlier response's part.
    """

    def __init__(self) -> None:
        self._next_index = 0

    def _allocate(self) -> int:
        index = self._next_index
        self._next_index += 1
        return index

    def remap_scoped(
        self, event: AgentBridgeEvent, stream_map: dict[int, int]
    ) -> AgentBridgeEvent:
        """Remap within one response stream via first-sight allocation.

        Allocation is synchronous on the event loop, so concurrent streams
        can never claim the same slot, while a repeated same-index start
        reuses its reservation and keeps replacement semantics.
        """
        local = event.part_index
        if local is None:
            return event
        global_index = stream_map.get(local)
        if global_index is None:
            global_index = self._allocate()
            stream_map[local] = global_index
        return event if global_index == local else replace(event, part_index=global_index)

    def remap_unscoped(
        self, event: AgentBridgeEvent, sticky_map: dict[int, int]
    ) -> AgentBridgeEvent:
        """Remap events whose stream boundaries are not observable.

        The v1 one-arg handler convention delivers bare events, so a part
        start at an already-seen local index is indistinguishable from a new
        response's first part. Every start therefore claims a fresh global
        part — degrading a true same-index restart to an appended part — and
        later deltas follow their local index's most recent allocation.
        """
        local = event.part_index
        if local is None:
            return event
        if event.kind == "text_replace" or local not in sticky_map:
            sticky_map[local] = self._allocate()
        global_index = sticky_map[local]
        return event if global_index == local else replace(event, part_index=global_index)


_DEFAULT_HISTORY_KEY = "__default__"
_PYDANTIC_AI_WARMUP_TIMEOUT_SECONDS = 5.0


class PydanticAIBridge:
    """Bridge for ``pydantic_ai.Agent`` and ``pydantic_graph.Graph``.

    Accepts either ``agent=`` (single-agent mode) or ``graph=`` with
    ``state_factory`` and ``initial_node_factory`` (graph mode).
    """

    COMMITTABLE_BOUNDARIES: ClassVar[Mapping[UnitKind | str, CommitRule]] = {
        UnitKind.AGENT: CommitRule.BETWEEN_TURNS,
        UnitKind.MODEL_NODE: CommitRule.NON_COMMITTABLE,
        UnitKind.TOOL_CALL: CommitRule.BETWEEN_PHASES,
        UnitKind.USER_PROMPT: CommitRule.NON_COMMITTABLE,
        UnitKind.WORKFLOW_NODE: CommitRule.BETWEEN_NODES,
    }

    def __init__(
        self,
        *,
        agent: Any | None = None,
        deps: Any = None,
        model_settings: Any = None,
        graph: Any | None = None,
        state_factory: Callable[[], Any] | None = None,
        initial_node_factory: Callable[[str, Any], Any] | None = None,
        agents: list[Any] | None = None,
        mcp_servers: list[Any] | None = None,
        toolsets: list[Any] | None = None,
    ) -> None:
        if mcp_servers is not None and toolsets is not None:
            raise BridgeInputError("Pass either mcp_servers= or toolsets= to PydanticAIBridge")
        if agent is not None and graph is not None:
            raise BridgeInputError("Cannot pass both agent= and graph= to PydanticAIBridge")
        if agent is None and graph is None:
            raise BridgeInputError("Must pass either agent= or graph= to PydanticAIBridge")

        self._agent: Any | None
        self._graph: Any | None
        self._state_factory: Callable[[], Any] | None
        self._initial_node_factory: Callable[[str, Any], Any] | None
        self._agents: list[Any] | None
        self._state: Any
        self._active_node: str | None

        if graph is not None:
            if state_factory is None or initial_node_factory is None:
                raise BridgeInputError(
                    "Graph mode requires state_factory and initial_node_factory"
                )
            # Convention validation: probe for _easycat_event_handler.
            probe = state_factory()
            if not hasattr(probe, "_easycat_event_handler"):
                raise BridgeConfigurationError(
                    "Graph state object must have an '_easycat_event_handler' "
                    "attribute. Add `_easycat_event_handler: Any = None` to "
                    "your state dataclass and pass "
                    "`event_stream_handler=ctx.state._easycat_event_handler` "
                    "to supported agent calls inside graph nodes (for "
                    "PydanticAI v2, use Agent.run() or Agent.run_stream(); "
                    "Agent.iter() does not accept event_stream_handler)."
                )
            self._mode = "graph"
            self._graph = graph
            self._state_factory = state_factory
            self._initial_node_factory = initial_node_factory
            self._agents = agents
            self._agent = None
            self._state = None
            self._active_node = None
        else:
            self._mode = "agent"
            self._agent = agent
            self._graph = None
            self._state_factory = None
            self._initial_node_factory = None
            self._agents = None
            self._state = None
            self._active_node = None

        self._deps = deps
        self._model_settings = model_settings
        self._mcp_servers = mcp_servers
        self._toolsets = toolsets
        self._message_histories: dict[str, list[Any]] = {_DEFAULT_HISTORY_KEY: []}
        self._graph_states: dict[str, Any] = {}
        self._last_history_key = _DEFAULT_HISTORY_KEY
        self._last_output: Any = None
        self._entered_warmup_agents: list[Any] = []
        self._warmup_lifecycle_lock = asyncio.Lock()
        self._warmup_closed = False

    async def warmup(self) -> None:
        """Enter static agents and prime their persistent model clients.

        PydanticAI otherwise creates and closes a deferred string model around
        each run. Resolving it once and keeping the agent context entered lets
        the first real turn reuse the connection established by the unbilled
        model-metadata request. Unsupported models and provider failures stay
        best-effort so warmup never becomes an availability gate.
        """
        async with self._warmup_lifecycle_lock:
            if self._warmup_closed:
                return
            agents = (self._agent,) if self._mode == "agent" else tuple(self._agents or ())
            for agent in agents:
                if agent is None or any(
                    entered is agent for entered in self._entered_warmup_agents
                ):
                    continue
                await self._warmup_agent(agent)

    async def _warmup_agent(self, agent: Any) -> None:
        try:
            model = self._resolve_warmup_model(agent)
        except Exception as exc:  # noqa: BLE001 intentional best-effort warmup boundary
            logger.debug("PydanticAI agent warmup skipped: %s", exc)
            return
        if not await self._enter_warmup_agent_context(agent):
            return
        await self._prime_warmup_model(model)

    @staticmethod
    def _resolve_warmup_model(agent: Any) -> Any:
        model = getattr(agent, "model", None)
        if isinstance(model, str):
            from pydantic_ai.models import infer_model

            model = infer_model(model)
            agent.model = model
        return model

    async def _enter_warmup_agent_context(self, agent: Any) -> bool:
        enter = getattr(agent, "__aenter__", None)
        exit_ = getattr(agent, "__aexit__", None)
        if enter is None or exit_ is None:
            return False
        try:
            async with asyncio.timeout(_PYDANTIC_AI_WARMUP_TIMEOUT_SECONDS):
                await enter()
        except asyncio.CancelledError as exc:
            await self._rollback_failed_warmup_entry(exit_, exc)
            raise
        except TimeoutError as exc:
            await self._rollback_failed_warmup_entry(exit_, exc)
            logger.debug("PydanticAI agent context warmup timed out")
            return False
        except Exception as exc:  # noqa: BLE001 intentional best-effort warmup boundary
            logger.debug("PydanticAI agent warmup skipped: %s", exc)
            return False
        self._entered_warmup_agents.append(agent)
        return True

    async def _prime_warmup_model(self, model: Any) -> None:
        client = getattr(model, "client", None)
        models = getattr(client, "models", None)
        retrieve = getattr(models, "retrieve", None)
        model_name = getattr(model, "model_name", None)
        if retrieve is None or not isinstance(model_name, str) or not model_name:
            return
        try:
            async with asyncio.timeout(_PYDANTIC_AI_WARMUP_TIMEOUT_SECONDS):
                result = retrieve(model_name)
                if inspect.isawaitable(result):
                    await result
        except Exception as exc:  # noqa: BLE001 intentional best-effort warmup boundary
            logger.debug("PydanticAI model warmup skipped: %s", exc)

    async def _rollback_failed_warmup_entry(self, exit_: Any, exc: BaseException) -> None:
        """Best-effort unwind a context whose bounded entry did not complete."""
        try:
            async with asyncio.timeout(_PYDANTIC_AI_WARMUP_TIMEOUT_SECONDS):
                await exit_(type(exc), exc, exc.__traceback__)
        except BaseException:
            logger.debug(
                "PydanticAI partial agent context rollback failed",
                exc_info=True,
            )

    async def rollback_warmup(self) -> None:
        """Release entered contexts after failed startup while allowing retry."""
        async with self._warmup_lifecycle_lock:
            await self._close_warmup_agents()

    async def aclose(self) -> None:
        """Exit agent contexts entered by :meth:`warmup`."""
        async with self._warmup_lifecycle_lock:
            self._warmup_closed = True
            await self._close_warmup_agents()

    async def _close_warmup_agents(self) -> None:
        """Exit and forget contexts entered by warmup while lifecycle is serialized."""
        entered, self._entered_warmup_agents = (
            self._entered_warmup_agents,
            [],
        )
        for agent in reversed(entered):
            try:
                await agent.__aexit__(None, None, None)
            except Exception:
                logger.debug(
                    "PydanticAI agent context close failed",
                    exc_info=True,
                )

    @property
    def _message_history(self) -> list[Any]:
        """Backward-compatible default-session message history view."""
        return self._history_for_key(_DEFAULT_HISTORY_KEY)

    @_message_history.setter
    def _message_history(self, value: list[Any]) -> None:
        self._set_history_for_key(_DEFAULT_HISTORY_KEY, value)
        self._last_history_key = _DEFAULT_HISTORY_KEY

    def _history_key_for_recorder(self, recorder: AgentRecorder | None) -> str:
        if recorder is None:
            return _DEFAULT_HISTORY_KEY
        session_id = getattr(getattr(recorder, "context", None), "session_id", "")
        if isinstance(session_id, str) and session_id.strip():
            return session_id
        return _DEFAULT_HISTORY_KEY

    def _history_for_key(self, history_key: str) -> list[Any]:
        return self._message_histories.setdefault(history_key, [])

    def _set_history_for_key(self, history_key: str, history: list[Any]) -> None:
        self._message_histories[history_key] = list(history)
        self._last_history_key = history_key

    def _graph_state_for_key(self, history_key: str) -> Any:
        state = self._graph_states.get(history_key)
        if state is None:
            state_factory = self._state_factory
            if state_factory is None:
                raise BridgeConfigurationError("PydanticAIBridge is in agent mode (no graph set)")
            state = state_factory()
            self._graph_states[history_key] = state
        self._state = state
        return state

    def _graph_state_for_key_if_present(self, history_key: str | None) -> Any | None:
        state = self._graph_states.get(history_key or self._last_history_key)
        if state is not None:
            self._state = state
        return state

    # ── ExternalAgentBridge interface ─────────────────────────────

    def version_info(self) -> dict[str, str]:
        """Return model and SDK metadata for journal reproducibility."""
        return bridge_version_info(
            provider="pydantic_ai",
            model=self._model_name(),
            distribution="pydantic-ai",
        )

    def _model_name(self) -> str | None:
        candidates = [self._agent, *(self._agents or [])]
        for agent in candidates:
            model = resolve_model_name(getattr(agent, "model", None))
            if model is not None:
                return model
        return None

    async def invoke(
        self,
        turn_input: AgentTurnInput,
        recorder: AgentRecorder,
        cancel_token: CancelToken | None = None,
    ) -> AsyncIterator[AgentBridgeEvent]:
        if self._mode == "graph":
            stream = self._invoke_graph(turn_input, recorder, cancel_token)
        else:
            stream = self._invoke_agent(turn_input, recorder, cancel_token)
        try:
            async for event in stream:
                yield event
        finally:
            # ``async for`` does not forward an early consumer ``aclose()``
            # (barge-in ``GeneratorExit``) into the delegated generator; close
            # it explicitly so ``turn_cursor``'s BaseException arm and the
            # mcp_servers restore run synchronously — before a follow-up
            # ``apply_interruption()`` — instead of at GC time.
            await aclose_quietly(stream)

    def snapshot_state(self) -> FrameworkStateSnapshot:
        if self._mode == "graph":
            fields: dict[str, Any] = {
                "mode": "graph",
                "graph": type(self._graph).__name__ if self._graph else "",
                "active_node": self._active_node,
                "session_count": len(self._graph_states),
            }
            # State may be large — use artifact ref if needed.
            state_ref = None
            if self._state is not None:
                try:
                    state_json = serialize_framework_state(_safe_state_dict(self._state))
                    if len(state_json) > 4096:
                        state_ref = f"state-{uuid4().hex[:16]}"
                        fields["state_summary"] = type(self._state).__name__
                    else:
                        fields["state"] = json.loads(state_json.decode("utf-8"))
                except (TypeError, ValueError, UnicodeDecodeError):
                    fields["state_summary"] = type(self._state).__name__
            return FrameworkStateSnapshot(
                fields=fields,
                state_ref=state_ref,
                kind="pydantic_ai_graph",
            )
        return FrameworkStateSnapshot(
            fields={
                "mode": "agent",
                "agent": _agent_name(self._agent),
                "turn_count": len(self._history_for_key(self._last_history_key)),
                "session_count": len(self._message_histories),
            },
            kind="pydantic_ai_agent",
        )

    def apply_interruption(
        self,
        delivered_text: str,
        mode: CancellationMode,
        recorder: AgentRecorder | None = None,
        caused_by_signal_id: str | None = None,
    ) -> None:
        # Step 1: plan the mutation.
        history_key = self._history_key_for_recorder(recorder)
        self._last_history_key = history_key
        plan = self._plan_interruption(delivered_text, mode, history_key)
        run_interruption_journal_protocol(
            plan,
            mode,
            recorder,
            caused_by_signal_id,
            serialize_state=lambda: self._serialize_framework_state(history_key),
            apply_mutation=lambda mutation_plan: self._apply_planned_mutation(
                mutation_plan, history_key
            ),
        )

    def _serialize_framework_state(self, history_key: str | None = None) -> bytes:
        """Serialize message history for artifact storage."""
        history = self._history_for_key(history_key or self._last_history_key)
        return serialize_framework_state(history, fallback=b"[]")

    def _plan_interruption(
        self, delivered_text: str, mode: CancellationMode, history_key: str | None = None
    ) -> InterruptionPlan:
        replacement = delivered_text + "..." if delivered_text else ""
        mutation_kind = "interrupt_truncate"
        if self._mode == "graph":
            mutation_kind = "interrupt_truncate_graph"
        history = self._history_for_key(history_key or self._last_history_key)
        state = self._graph_state_for_key_if_present(history_key)
        pre_ref = f"pydantic-pre-{id(history):x}"
        post_ref = f"pydantic-post-{id(history):x}"
        return InterruptionPlan(
            mutation_kind=mutation_kind,
            pre_state_ref=pre_ref,
            post_state_ref=post_ref,
            framework_instructions={
                "replacement": replacement,
                "mode": self._mode,
                "use_custom_truncation": (
                    self._mode == "graph"
                    and state is not None
                    and hasattr(state, "truncate_last_assistant")
                ),
                "delivered_text": delivered_text,
            },
        )

    def _apply_planned_mutation(
        self, plan: InterruptionPlan, history_key: str | None = None
    ) -> None:
        instructions = plan.framework_instructions
        replacement = instructions["replacement"]

        if instructions.get("use_custom_truncation"):
            state = self._graph_state_for_key_if_present(history_key)
            if state is not None:
                state.truncate_last_assistant(instructions["delivered_text"])
            return

        try:
            from pydantic_ai.messages import (
                ModelRequest,
                ModelResponse,
                TextPart,
                UserPromptPart,
            )
        except ImportError:
            return

        history = self._history_for_key(history_key or self._last_history_key)
        for i in range(len(history) - 1, -1, -1):
            msg = history[i]
            if isinstance(msg, ModelRequest) and any(
                isinstance(part, UserPromptPart) for part in msg.parts
            ):
                # PydanticAI may expose the current user request before it
                # commits a streaming ModelResponse. Preserve what was heard
                # as this turn's response rather than walking across the user
                # boundary and rewriting the previous completed turn.
                if replacement:
                    history.append(ModelResponse(parts=[TextPart(content=replacement)]))
                break
            if isinstance(msg, ModelResponse):
                text_parts: list[Any] = [
                    part for part in msg.parts if type(part).__name__ == "TextPart"
                ]
                for idx, part in enumerate(text_parts):
                    value = replacement if idx == 0 else ""
                    try:
                        part.content = value
                    except (AttributeError, TypeError):
                        object.__setattr__(part, "content", value)
                break

    def reset(self) -> None:
        for history in self._message_histories.values():
            history.clear()
        self._message_histories = {_DEFAULT_HISTORY_KEY: []}
        self._graph_states.clear()
        self._last_history_key = _DEFAULT_HISTORY_KEY
        self._last_output = None
        self._state = None
        self._active_node = None

    def configure_runtime(
        self,
        *,
        mcp_servers: list[str] | None = None,
        model: str | None = None,
        api_key: str | None = None,
    ) -> None:
        """Apply session-level MCP servers (model/api_key are unused here)."""
        if mcp_servers is not None:
            self._mcp_servers = list(mcp_servers)

    # ── History post-processing ───────────────────────────────────

    def replace_last_assistant_text(self, text: str) -> None:
        """Rewrite the last assistant ``TextPart`` in history to ``text``.

        Called by the adapter shim after post-processing (e.g. Markdown
        stripping) so that subsequent turns condition on the cleaned text
        rather than the raw LLM output.
        """
        try:
            from pydantic_ai.messages import ModelRequest, ModelResponse, UserPromptPart
        except ImportError:
            return
        for msg in reversed(self._history_for_key(self._last_history_key)):
            if isinstance(msg, ModelRequest) and any(
                isinstance(part, UserPromptPart) for part in msg.parts
            ):
                return
            if not isinstance(msg, ModelResponse):
                continue
            text_parts: list[Any] = [
                part for part in msg.parts if type(part).__name__ == "TextPart"
            ]
            if not text_parts:
                return
            originals = [p.content for p in text_parts]
            replacements = split_replacement_by_original_parts(originals, text)
            for part, repl in zip(text_parts, replacements):
                try:
                    part.content = repl
                except (AttributeError, TypeError):
                    object.__setattr__(part, "content", repl)
            return

    def append_interruption_note(self, note: str) -> None:
        """Append an interruption note to PydanticAI message history.

        Appends a ``ModelRequest`` containing a ``SystemPromptPart`` so
        the next turn sees the interruption signal. Best effort — if the
        required message types can't be imported, this is a no-op.
        """
        try:
            from pydantic_ai.messages import ModelRequest, SystemPromptPart
        except ImportError:
            return
        try:
            self._history_for_key(self._last_history_key).append(
                ModelRequest(parts=[SystemPromptPart(content=note)])
            )
        except Exception:
            logger.debug("Failed to append interruption note to PydanticAI history", exc_info=True)

    # ── Agent mode ───────────────────────────────────────────────

    def _require_agent(self) -> Any:
        """Return the wrapped agent; only callable in agent mode."""
        agent = self._agent
        if agent is None:
            raise BridgeConfigurationError("PydanticAIBridge is in graph mode (no agent set)")
        return agent

    async def _invoke_agent(
        self,
        turn_input: AgentTurnInput,
        recorder: AgentRecorder,
        cancel_token: CancelToken | None,
    ) -> AsyncIterator[AgentBridgeEvent]:
        history_key = self._history_key_for_recorder(recorder)
        self._last_history_key = history_key
        agent = self._require_agent()

        agent_cursor = ExecutionCursor(
            unit_id=f"agent-{uuid4().hex[:8]}",
            unit_kind=UnitKind.AGENT,
            display_name=_agent_name(agent),
            entered_at=time.monotonic_ns(),
            committable=False,
        )
        accumulated = AgentTextStream()
        raw_output: Any = None
        done_emitted = False

        # ``turn_cursor`` centralizes enter → error → BaseException → clean
        # exit; the inner bare ``try/finally`` restores the agent's mcp_servers
        # on every path (including cancellation) before the cm records the exit.
        with recorder.turn_cursor(agent_cursor):
            saved_mcp_servers = _UNSET
            inner = None
            try:
                if self._mcp_servers is not None and hasattr(agent, "mcp_servers"):
                    saved_mcp_servers = getattr(agent, "mcp_servers", None)
                    agent.mcp_servers = list(self._mcp_servers)
                if hasattr(agent, "iter"):
                    inner = self._stream_via_iter(turn_input, recorder, cancel_token, history_key)
                    async for ev in inner:
                        if ev.kind == "done":
                            done_emitted = True
                        else:
                            accumulated.apply(ev)
                        yield ev
                    raw_output = self._last_output
                else:
                    inner = self._stream_via_run_stream(
                        turn_input,
                        recorder,
                        cancel_token,
                        history_key,
                    )
                    async for ev in inner:
                        accumulated.apply(ev)
                        yield ev
                    raw_output = self._last_output
            finally:
                # Close the delegated stream first: ``_stream_via_iter`` holds
                # the ``agent.iter()`` context open, and ``async for`` does not
                # forward a barge-in ``GeneratorExit`` into it — without this
                # its teardown would only run at GC time.
                if inner is not None:
                    await aclose_quietly(inner)
                if saved_mcp_servers is not _UNSET:
                    agent.mcp_servers = saved_mcp_servers

        if not done_emitted:
            yield AgentBridgeEvent(
                kind="done",
                text=_completion_text(accumulated.text, raw_output),
                structured_output=raw_output,
            )

    async def _stream_via_iter(
        self,
        turn_input: AgentTurnInput,
        recorder: AgentRecorder,
        cancel_token: CancelToken | None,
        history_key: str,
    ) -> AsyncIterator[AgentBridgeEvent]:
        """Stream using ``agent.iter()`` with full event capture."""
        agent = self._require_agent()
        tool_call_ids: dict[int, str] = {}
        part_indexes = _RunPartIndexAllocator()
        async with agent.iter(
            turn_input.text,
            **self._agent_run_kwargs(agent.iter, turn_input, history_key),
        ) as agent_run:
            try:
                interrupted = False
                async for node in agent_run:
                    node_cls = type(node).__name__
                    is_tool_node = node_cls == "CallToolsNode"

                    if cancel_token and cancel_token.is_cancelled:
                        interrupted = True
                        if not is_tool_node:
                            break

                    if not hasattr(node, "stream"):
                        continue

                    async with node.stream(agent_run.ctx) as stream:
                        # Each node streams one model response whose part
                        # indexes restart at 0; namespace them per stream.
                        node_part_index_map: dict[int, int] = {}
                        async for event in stream:
                            if cancel_token and cancel_token.is_cancelled and not interrupted:
                                interrupted = True
                                if not is_tool_node:
                                    break
                            mapped = translate_event(
                                event,
                                recorder,
                                tool_call_ids=tool_call_ids,
                            )
                            if mapped is not None:
                                if interrupted and is_text_event(mapped):
                                    continue
                                mapped = part_indexes.remap_scoped(mapped, node_part_index_map)
                                yield mapped

            except GeneratorExit:
                # Do not inject GeneratorExit into PydanticAI's Agent.iter
                # context: v1 records it as a node failure and re-raises it
                # during __aexit__. Snapshot the partial turn, then consume the
                # normal async-generator close signal here so the SDK context
                # exits cleanly.
                await self._commit_interrupted_agent_run(history_key, agent_run)
                return
            except BaseException:
                # A hard task cancel or consumer ``aclose()`` can unwind at the
                # suspended yield above before normal fall-through commits this
                # turn. PydanticAI exposes ``new_messages()`` while the run is
                # still active, so snapshot it before the context exits; the
                # follow-up interruption rewrite must target this turn rather
                # than the previous completed assistant response.
                await self._commit_interrupted_agent_run(history_key, agent_run)
                raise
            else:
                await self._commit_agent_run(history_key, agent_run)
            finally:
                await record_usage_from_result(
                    recorder,
                    agent_run,
                    provider="pydantic_ai",
                    model=self._model_name(),
                )

    async def _stream_via_run_stream(
        self,
        turn_input: AgentTurnInput,
        recorder: AgentRecorder,
        cancel_token: CancelToken | None,
        history_key: str,
    ) -> AsyncIterator[AgentBridgeEvent]:
        """Fallback: stream text-only via ``agent.run_stream()``."""
        agent = self._require_agent()
        async with agent.run_stream(
            turn_input.text,
            **self._agent_run_kwargs(agent.run_stream, turn_input, history_key),
        ) as result:
            try:
                accumulated = ""
                async for full_text in result.stream_text():
                    if cancel_token and cancel_token.is_cancelled:
                        break
                    delta = full_text[len(accumulated) :]
                    if delta:
                        yield AgentBridgeEvent(kind="text_delta", text=delta)
                    accumulated = full_text

            except GeneratorExit:
                await self._commit_interrupted_agent_run(history_key, result)
                return
            except BaseException:
                await self._commit_interrupted_agent_run(history_key, result)
                raise
            else:
                await self._commit_agent_run(history_key, result)
            finally:
                await record_usage_from_result(
                    recorder,
                    result,
                    provider="pydantic_ai",
                    model=self._model_name(),
                )

    async def _commit_interrupted_agent_run(self, history_key: str, run: Any) -> None:
        """Best-effort history snapshot while cancellation is already active."""
        try:
            await self._commit_agent_run(history_key, run)
        except (asyncio.CancelledError, Exception):
            # Never replace the cancellation/GeneratorExit with a secondary
            # SDK snapshot failure. The user-boundary guard below still keeps
            # a missing current turn from clobbering the prior assistant.
            logger.warning(
                "Failed to snapshot interrupted PydanticAI run history",
                exc_info=True,
            )

    async def _commit_agent_run(self, history_key: str, run: Any) -> None:
        """Append the SDK's current-turn messages to bridge-owned history."""
        new_messages = await _run_new_messages(run)
        history = self._history_for_key(history_key)
        self._set_history_for_key(history_key, [*history, *new_messages])
        self._last_output = await _run_output(run)

    def _runtime_toolsets(self) -> list[Any] | None:
        if self._toolsets is not None:
            return list(self._toolsets)
        if self._mcp_servers is not None:
            return _mcp_servers_to_toolsets(self._mcp_servers)
        return None

    def _agent_run_kwargs(
        self, call: Callable[..., Any], turn_input: AgentTurnInput, history_key: str | None = None
    ) -> dict[str, Any]:
        kwargs: dict[str, Any] = {
            "message_history": self._message_history_for_turn(turn_input.context, history_key),
            "deps": self._deps,
            "model_settings": self._model_settings,
        }
        if _supports_kwarg(call, "toolsets"):
            runtime_toolsets = self._runtime_toolsets()
            if runtime_toolsets is not None:
                kwargs["toolsets"] = list(runtime_toolsets)
        return kwargs

    def _message_history_for_turn(
        self, context: list[dict[str, str]], history_key: str | None = None
    ) -> list[Any] | None:
        history = list(self._history_for_key(history_key or self._last_history_key))
        context_messages = normalize_context_messages(context, own_history=bool(history))
        if not context_messages:
            return history or None
        try:
            converted = [_pydantic_message_from_context(item) for item in context_messages]
        except ImportError:
            logger.debug("PydanticAI message types unavailable; dropping turn context")
            return history or None
        return [*converted, *history]

    def _graph_iter(self, initial_input: Any, state: Any) -> Any:
        graph = self._graph
        if graph is None:
            raise BridgeConfigurationError("PydanticAIBridge is in agent mode (no graph set)")
        graph_iter = graph.iter
        kwargs: dict[str, Any] = {"state": state}
        if self._deps is not None and _supports_kwarg(graph_iter, "deps"):
            kwargs["deps"] = self._deps
        if _supports_kwarg(graph_iter, "inputs"):
            kwargs["inputs"] = initial_input
        elif _supports_kwarg(graph_iter, "input"):
            kwargs["input"] = initial_input
        elif _supports_positional_arg(graph_iter):
            return graph_iter(initial_input, state=state)
        else:
            raise BridgeConfigurationError(
                "PydanticAIBridge Graph mode could not determine how to pass graph input. "
                "Expected graph.iter(initial_node, state=...) or keyword "
                "graph.iter(state=..., inputs=...)."
            )
        return graph_iter(**kwargs)

    # ── Graph mode ───────────────────────────────────────────────

    async def _invoke_graph(
        self,
        turn_input: AgentTurnInput,
        recorder: AgentRecorder,
        cancel_token: CancelToken | None,
    ) -> AsyncIterator[AgentBridgeEvent]:
        history_key = self._history_key_for_recorder(recorder)
        self._last_history_key = history_key

        initial_node_factory = self._initial_node_factory
        if initial_node_factory is None:
            raise BridgeConfigurationError("PydanticAIBridge is in agent mode (no graph set)")
        state = self._graph_state_for_key(history_key)
        initial_node = initial_node_factory(turn_input.text, state)

        # Install event handler on state for per-agent capture.
        _handler = _GraphEventHandler(recorder)
        state._easycat_event_handler = _handler

        accumulated = ""
        graph_output: Any = None
        prev_node_name: str | None = None
        prev_cursor: ExecutionCursor | None = None

        # Inject MCP servers/toolsets into graph agents (same save/restore pattern).
        saved_mcp: list[tuple[Any, Any]] = []
        graph_agent_overrides = ExitStack()
        observed_history: list[Any] = []

        try:
            if self._agents:
                runtime_toolsets: list[Any] | None = None
                for ag in self._agents:
                    if self._mcp_servers is not None and hasattr(ag, "mcp_servers"):
                        saved_mcp.append((ag, getattr(ag, "mcp_servers", None)))
                        ag.mcp_servers = list(self._mcp_servers)
                    elif hasattr(ag, "override") and _supports_kwarg(ag.override, "toolsets"):
                        if runtime_toolsets is None:
                            runtime_toolsets = self._runtime_toolsets()
                        if runtime_toolsets is not None:
                            override = ag.override(toolsets=list(runtime_toolsets))
                            graph_agent_overrides.enter_context(override)

            with graph_agent_overrides:
                async with self._graph_iter(initial_node, state) as graph_run:
                    async for graph_item in graph_run:
                        graph_units = _graph_iteration_units(graph_item)
                        if not graph_units:
                            output = _graph_item_output(graph_item)
                            if output is not _UNSET:
                                graph_output = output
                                self._last_output = output
                                if isinstance(output, str):
                                    accumulated = accumulated or output
                            continue

                        for node in graph_units:
                            node_name = _graph_item_display_name(node)
                            observed_history.append(node)
                            self._active_node = node_name

                            if cancel_token and cancel_token.is_cancelled:
                                break

                            # Emit handoff triple if node changed.
                            if prev_node_name is not None and node_name != prev_node_name:
                                for ev in _handler.drain():
                                    yield ev
                                if prev_cursor is not None:
                                    recorder.record_unit_exited(
                                        prev_cursor.with_committable(True), reason=None
                                    )
                                recorder.record_framework_handoff(
                                    from_unit=prev_node_name,
                                    to_unit=node_name,
                                    reason="graph_transition",
                                )

                            cursor = ExecutionCursor(
                                unit_id=f"node-{uuid4().hex[:8]}",
                                unit_kind=UnitKind.WORKFLOW_NODE,
                                display_name=node_name,
                                entered_at=time.monotonic_ns(),
                                committable=False,
                            )
                            recorder.record_unit_entered(cursor)
                            prev_cursor = cursor
                            prev_node_name = node_name

                            # Forward any events the handler captured during this node.
                            for ev in _handler.drain():
                                yield ev

                        if cancel_token and cancel_token.is_cancelled:
                            break

            for ev in _handler.drain():
                yield ev

            # Close the last cursor.
            if prev_cursor is not None:
                recorder.record_unit_exited(prev_cursor.with_committable(True), reason=None)
                prev_cursor = None

            # Collect any text from the handler.
            accumulated = _handler.accumulated_text

            # Capture result.
            output = await _run_output(graph_run)
            if output is not None:
                graph_output = output
                self._last_output = output
                if isinstance(output, str):
                    accumulated = accumulated or output
            # Capture graph run history as a state snapshot record.
            run_history = getattr(graph_run, "history", None)
            if run_history is None:
                run_history = observed_history
            if run_history is not None:
                try:
                    serialized = json.dumps(
                        [
                            {"node": type(n).__name__, "index": i}
                            for i, n in enumerate(run_history)
                        ],
                        default=str,
                    )
                    recorder.record_state_snapshot(
                        ref=f"graph-history-{uuid4().hex[:8]}",
                        payload=serialized.encode("utf-8"),
                    )
                except (TypeError, ValueError):
                    logger.debug("Failed to serialize graph run history")

            # Populate _message_history so the default interruption
            # truncation path has something to work with in graph mode.
            if accumulated:
                try:
                    from pydantic_ai.messages import ModelResponse, TextPart

                    self._set_history_for_key(
                        history_key, [ModelResponse(parts=[TextPart(content=accumulated)])]
                    )
                except ImportError:
                    pass

            # Convention enforcement: verify handler was actually used.
            if not _handler.was_called and prev_node_name is not None:
                from easycat.integrations.agents.base import (
                    ConventionViolationError,
                )

                raise ConventionViolationError(
                    "PydanticAIBridge Graph mode: the _easycat_event_handler "
                    "was installed on the state but no agent call passed it "
                    "through via event_stream_handler=. Ensure every supported "
                    "agent call inside graph nodes passes "
                    "event_stream_handler=ctx.state._easycat_event_handler. "
                    "In PydanticAI v2, Agent.iter() does not accept "
                    "event_stream_handler; use Agent.run(), Agent.run_stream(), "
                    "or configure an agent-level event stream handler instead."
                )

        except Exception as exc:
            recorder.record_framework_error(ErrorInfo.from_exception(exc))
            if prev_cursor is not None:
                recorder.record_unit_exited(prev_cursor, reason="error")
            raise
        except BaseException:
            # ``AgentRunner``'s timeout/barge-in ``aclose()`` injects
            # ``asyncio.CancelledError`` / ``GeneratorExit`` here; neither
            # is an ``Exception`` so the block above is skipped and the
            # still-open graph-node cursor would be left without a
            # ``unit_exited`` record, breaking the recorder's strict stack
            # invariant for the postmortem journal.  Close it defensively
            # before re-raising; no ``record_framework_error`` since a
            # cancelled turn isn't a framework fault.
            if prev_cursor is not None:
                recorder.safe_exit_cursor(prev_cursor)
            raise
        finally:
            graph_agent_overrides.close()
            for ag, original in saved_mcp:
                ag.mcp_servers = original

        yield AgentBridgeEvent(
            kind="done",
            text=_completion_text(accumulated, graph_output),
            structured_output=graph_output,
        )


# ── Graph event handler ──────────────────────────────────────────


class _GraphEventHandler:
    """Installed on graph state to capture per-agent events.

    Agents inside graph nodes call this via
    ``event_stream_handler=ctx.state._easycat_event_handler``.
    """

    def __init__(self, recorder: AgentRecorder) -> None:
        self._recorder = recorder
        self._text_stream = AgentTextStream()
        self._was_called = False
        self._pending: list[AgentBridgeEvent] = []
        self._tool_call_ids: dict[int, str] = {}
        self._part_indexes = _RunPartIndexAllocator()
        # v1 single events carry no stream identity; this sticky map routes
        # each local index's deltas to its most recent part allocation.
        self._v1_part_index_map: dict[int, int] = {}

    @property
    def accumulated_text(self) -> str:
        return self._text_stream.text

    @property
    def was_called(self) -> bool:
        return self._was_called

    def drain(self) -> list[AgentBridgeEvent]:
        events = self._pending
        self._pending = []
        return events

    async def _handle_event(
        self,
        event: Any,
        *,
        part_index_map: dict[int, int] | None = None,
        tool_call_ids: dict[int, str] | None = None,
    ) -> None:
        mapped = translate_event(
            event,
            self._recorder,
            tool_call_ids=self._tool_call_ids if tool_call_ids is None else tool_call_ids,
        )
        if mapped is not None:
            if part_index_map is not None:
                mapped = self._part_indexes.remap_scoped(mapped, part_index_map)
            else:
                mapped = self._part_indexes.remap_unscoped(mapped, self._v1_part_index_map)
            self._text_stream.apply(mapped)
            self._pending.append(mapped)

    async def __call__(self, *args: Any) -> None:
        """Handle PydanticAI v1 single events and v2 ``(ctx, events)`` streams."""
        self._was_called = True
        if len(args) == 1:
            await self._handle_event(args[0])
            return
        if len(args) == 2:
            ctx, events = args
            part_index_map: dict[int, int] = {}
            tool_call_ids: dict[int, str] = {}
            try:
                async for event in events:
                    await self._handle_event(
                        event,
                        part_index_map=part_index_map,
                        tool_call_ids=tool_call_ids,
                    )
            finally:
                await record_usage_from_result(
                    self._recorder,
                    ctx,
                    provider="pydantic_ai",
                    model=resolve_model_name(getattr(ctx, "model", None)),
                )
            return
        raise TypeError("_GraphEventHandler expects event or (ctx, events)")


_UNSET = object()


def _supports_positional_arg(call: Callable[..., Any]) -> bool:
    try:
        signature = inspect.signature(call)
    except (TypeError, ValueError):
        return True
    for param in signature.parameters.values():
        if param.kind == inspect.Parameter.VAR_POSITIONAL:
            return True
        if param.kind in (
            inspect.Parameter.POSITIONAL_ONLY,
            inspect.Parameter.POSITIONAL_OR_KEYWORD,
        ):
            return param.name != "self"
    return False


def _supports_kwarg(call: Callable[..., Any], name: str) -> bool:
    try:
        signature = inspect.signature(call)
    except (TypeError, ValueError):
        return False
    for param in signature.parameters.values():
        if param.kind == inspect.Parameter.VAR_KEYWORD:
            return True
    return name in signature.parameters


def _mcp_servers_to_toolsets(mcp_servers: list[Any]) -> list[Any]:
    try:
        from pydantic_ai.mcp import MCPToolset, StdioTransport
    except ImportError as exc:
        raise BridgeConfigurationError(
            "PydanticAIBridge mcp_servers= requires PydanticAI v2 MCP support. "
            "Install the pydantic-ai MCP extra or pass explicit toolsets= objects."
        ) from exc

    toolsets: list[Any] = []
    for index, server in enumerate(mcp_servers):
        if not isinstance(server, str):
            raise BridgeConfigurationError(
                "PydanticAIBridge mcp_servers= expects EasyCat MCP URI strings. "
                "Pass PydanticAI toolset objects with toolsets= instead."
            )

        if server.startswith("stdio://"):
            command_line = unquote(server[len("stdio://") :]).strip()
            try:
                parts = shlex.split(command_line)
            except ValueError as exc:
                raise BridgeConfigurationError(
                    f"Invalid stdio MCP server URI at index {index}: {server!r}"
                ) from exc
            if not parts:
                raise BridgeConfigurationError(
                    f"Invalid stdio MCP server URI at index {index}: missing command"
                )
            toolsets.append(MCPToolset(StdioTransport(command=parts[0], args=parts[1:])))
            continue

        if server.startswith(("http://", "https://")):
            toolsets.append(MCPToolset(server))
            continue

        if server.startswith("sse://"):
            raise BridgeConfigurationError(
                "PydanticAI v2 MCPToolset does not accept EasyCat sse:// MCP URIs. "
                "Use an http(s):// SSE endpoint such as http://host/sse, or pass "
                "a pre-built PydanticAI MCPToolset with toolsets=."
            )

        raise BridgeConfigurationError(
            f"Unsupported MCP server URI for PydanticAIBridge: {server!r}. "
            "Use stdio://, http://, or https://, or pass explicit toolsets= objects."
        )

    return toolsets


def _graph_iteration_units(item: Any) -> list[Any]:
    if _is_graph_end_item(item):
        return []
    if isinstance(item, Sequence) and not isinstance(item, (str, bytes, bytearray)):
        return list(item)
    return [item]


def _graph_item_output(item: Any) -> Any:
    if not _is_graph_end_item(item):
        return _UNSET
    if hasattr(item, "value"):
        return item.value
    return getattr(item, "_value", None)


def _is_graph_end_item(item: Any) -> bool:
    return type(item).__name__ == "EndMarker"


def _graph_item_display_name(item: Any) -> str:
    node_id = getattr(item, "node_id", None)
    if node_id is not None:
        node_id_text = str(node_id)
        if node_id_text != "__start__":
            return node_id_text

    inputs = getattr(item, "inputs", _UNSET)
    if inputs is not _UNSET and not isinstance(
        inputs, (str, bytes, bytearray, int, float, bool, type(None))
    ):
        return type(inputs).__name__

    return type(item).__name__


async def _maybe_await(value: Any) -> Any:
    if inspect.isawaitable(value):
        return await value
    return value


async def _run_output(run: Any) -> Any:
    output = getattr(run, "output", None)
    output = await _maybe_await(output)
    if output is not None:
        return output

    result = getattr(run, "result", None)
    result = await _maybe_await(result)
    if result is not None:
        output = getattr(result, "output", None)
        output = await _maybe_await(output)
        if output is not None:
            return output

    get_output = getattr(run, "get_output", None)
    if callable(get_output):
        try:
            return await _maybe_await(get_output())
        except Exception:
            logger.debug("Failed to read PydanticAI run output", exc_info=True)
    return None


async def _run_new_messages(run: Any) -> list[Any]:
    new_messages = getattr(run, "new_messages", None)
    if callable(new_messages):
        return list(await _maybe_await(new_messages()))

    result = getattr(run, "result", None)
    result = await _maybe_await(result)
    if result is not None:
        result_new_messages = getattr(result, "new_messages", None)
        if callable(result_new_messages):
            return list(await _maybe_await(result_new_messages()))

    logger.debug("PydanticAI run object does not expose new_messages()")
    return []


# ── Helpers ──────────────────────────────────────────────────────


def _agent_name(agent: Any) -> str:
    return getattr(agent, "name", None) or type(agent).__name__


def _pydantic_message_from_context(item: dict[str, str]) -> Any:
    from pydantic_ai.messages import ModelRequest, ModelResponse, SystemPromptPart, TextPart

    role = item["role"]
    content = item["content"]
    if role == "assistant":
        return ModelResponse(parts=[TextPart(content=content)])
    if role in {"system", "developer"}:
        return ModelRequest(parts=[SystemPromptPart(content=content)])

    from pydantic_ai.messages import UserPromptPart

    return ModelRequest(parts=[UserPromptPart(content=content)])


def _safe_state_dict(state: Any) -> dict[str, Any]:
    """Extract public state fields without invoking value repr/string hooks."""
    from dataclasses import fields as dc_fields

    if isinstance(state, Mapping):
        return dict(state)
    result: dict[str, Any] = {}
    try:
        for f in dc_fields(state):
            if f.name.startswith("_"):
                continue
            result[f.name] = getattr(state, f.name, None)
    except TypeError:
        # Not a dataclass — use vars.
        try:
            state_vars = vars(state)
        except TypeError:
            return {}
        for k, v in state_vars.items():
            if k.startswith("_"):
                continue
            result[k] = v
    return result
