"""Async Agent framework for Prompture.

Provides :class:`AsyncAgent`, the async counterpart of :class:`~prompture.agent.Agent`.
All methods are ``async`` and use :class:`~prompture.async_conversation.AsyncConversation`.

Example::

    from prompture import AsyncAgent

    agent = AsyncAgent("openai/gpt-4o", system_prompt="You are helpful.")
    result = await agent.run("What is 2 + 2?")
    print(result.output)
"""

from __future__ import annotations

import asyncio
import inspect
import json
import logging
import time
import typing
from collections.abc import AsyncGenerator, Callable
from typing import Any, Generic

from pydantic import BaseModel

from ..extraction.tools import clean_json_text
from ..infra.budget import BudgetPolicy, BudgetState, enforce_budget
from ..infra.callbacks import DriverCallbacks
from ..infra.session import UsageSession
from .persona import Persona
from .tools_schema import ToolDefinition, ToolRegistry
from .types import (
    AgentCallbacks,
    AgentResult,
    AgentState,
    AgentStep,
    ApprovalRequired,
    DepsType,
    ModelRetry,
    RunContext,
    StepType,
    StreamEvent,
    StreamEventType,
)

logger = logging.getLogger("prompture.async_agent")

_OUTPUT_PARSE_MAX_RETRIES = 3
_OUTPUT_GUARDRAIL_MAX_RETRIES = 3
_DEFAULT_MAX_AGENT_DEPTH = 5

# Share the same ContextVar with the sync Agent so depth tracking crosses
# agent boundaries (e.g. a sync Agent calling an AsyncAgent as a tool).
from .agent import _agent_depth

# ------------------------------------------------------------------
# Helpers
# ------------------------------------------------------------------


def _is_async_callable(fn: Callable[..., Any]) -> bool:
    """Check if *fn* is an async callable (coroutine function or has async ``__call__``)."""
    if asyncio.iscoroutinefunction(fn):
        return True
    # Check if the object has an async __call__ method (callable class)
    dunder_call = type(fn).__call__ if callable(fn) else None
    return dunder_call is not None and asyncio.iscoroutinefunction(dunder_call)


def _tool_wants_context(fn: Callable[..., Any]) -> bool:
    """Check whether *fn*'s first parameter is annotated as :class:`RunContext`."""
    sig = inspect.signature(fn)
    params = list(sig.parameters.keys())
    if not params:
        return False

    first_param = params[0]
    if first_param == "self":
        if len(params) < 2:
            return False
        first_param = params[1]

    annotation = None
    try:
        hints = typing.get_type_hints(fn, include_extras=True)
        annotation = hints.get(first_param)
    except Exception:
        # get_type_hints can fail with local/forward references; fall back to raw annotation
        pass

    if annotation is None:
        raw = sig.parameters[first_param].annotation
        if raw is inspect.Parameter.empty:
            return False
        annotation = raw

    if isinstance(annotation, str):
        return annotation == "RunContext" or annotation.startswith("RunContext[")

    if annotation is RunContext:
        return True

    origin = getattr(annotation, "__origin__", None)
    return origin is RunContext


def _get_first_param_name(fn: Callable[..., Any]) -> str:
    """Return the name of the first non-self parameter of *fn*."""
    sig = inspect.signature(fn)
    for name, _param in sig.parameters.items():
        if name != "self":
            return name
    return ""


# ------------------------------------------------------------------
# AsyncAgent
# ------------------------------------------------------------------


class AsyncAgent(Generic[DepsType]):
    """Async agent that executes a ReAct loop with tool support.

    Mirrors :class:`~prompture.agent.Agent` but uses
    :class:`~prompture.async_conversation.AsyncConversation` and
    ``async`` methods throughout.

    Args:
        model: Model string in ``"provider/model"`` format.
        driver: Pre-built async driver instance.
        tools: Initial tools as a list of callables or a :class:`ToolRegistry`.
        system_prompt: System prompt prepended to every run.  May also be a
            callable ``(RunContext) -> str`` for dynamic prompts.
        output_type: Optional Pydantic model class for structured output.
        max_iterations: Maximum tool-use rounds per run.
        max_cost: Soft budget in USD.
        options: Extra driver options forwarded to every LLM call.
            Common keys include ``"temperature"`` (float, 0.0-2.0),
            ``"max_tokens"`` (int), and ``"top_p"`` (float).
            Example: ``options={"temperature": 0.7, "max_tokens": 1024}``.
        deps_type: Type hint for dependencies.
        agent_callbacks: Agent-level observability callbacks.
        input_guardrails: Functions called before the prompt is sent.
        output_guardrails: Functions called after output is parsed.
        persistent_conversation: When ``True``, subsequent ``run()``
            calls reuse the same :class:`AsyncConversation` so the model
            sees the full multi-turn history.  Default ``False``.
        security_context: Optional tukuy ``SecurityContext`` activated
            around tool execution so tukuy skills run with filesystem
            and network scoping.  When ``None`` (default) no scoping
            is applied.
        auto_approve_safe_only: When ``True``, tools backed by tukuy
            skills that declare ``side_effects=True`` or
            ``requires_network=True`` will raise
            :class:`ApprovalRequired` before execution.  Default
            ``False``.
        skill_config: Optional configuration dict injected as a
            :class:`SkillContext` into tukuy skill ``invoke()`` calls.
            When ``None`` (default) no config is injected.
    """

    def __init__(
        self,
        model: str = "",
        *,
        driver: Any | None = None,
        tools: list[Callable[..., Any]] | ToolRegistry | None = None,
        system_prompt: str | Persona | Callable[..., str] | None = None,
        output_type: type[BaseModel] | None = None,
        max_iterations: int = 10,
        max_cost: float | None = None,
        max_tokens: int | None = None,
        budget_policy: BudgetPolicy | None = None,
        fallback_models: list[str] | None = None,
        on_model_fallback: Callable[..., Any] | None = None,
        options: dict[str, Any] | None = None,
        deps_type: type | None = None,
        agent_callbacks: AgentCallbacks | None = None,
        input_guardrails: list[Callable[..., Any]] | None = None,
        output_guardrails: list[Callable[..., Any]] | None = None,
        name: str = "",
        description: str = "",
        output_key: str | None = None,
        persistent_conversation: bool = False,
        security_context: Any | None = None,
        auto_approve_safe_only: bool = False,
        skill_config: dict[str, Any] | None = None,
        max_tool_result_length: int | None = None,
        max_depth: int = _DEFAULT_MAX_AGENT_DEPTH,
    ) -> None:
        if not model and driver is None:
            raise ValueError("Either model or driver must be provided")

        self._model = model
        self._driver = driver
        self._max_depth = max_depth
        self._system_prompt = system_prompt
        self._output_type = output_type
        self._max_iterations = max_iterations
        self._max_cost = max_cost
        self._max_tokens = max_tokens
        self._budget_policy = budget_policy
        self._fallback_models = list(fallback_models) if fallback_models else None
        self._on_model_fallback = on_model_fallback
        self._options = dict(options) if options else {}
        self._deps_type = deps_type
        self._agent_callbacks = agent_callbacks or AgentCallbacks()
        self._input_guardrails = list(input_guardrails) if input_guardrails else []
        self._output_guardrails = list(output_guardrails) if output_guardrails else []
        self.name = name
        self.description = description
        self.output_key = output_key
        self._persistent_conversation = persistent_conversation
        self._security_context = security_context
        self._auto_approve_safe_only = auto_approve_safe_only
        self._skill_config = skill_config
        self._max_tool_result_length = max_tool_result_length
        self._conversation: Any = None

        # Build internal tool registry
        self._tools = ToolRegistry()
        if isinstance(tools, ToolRegistry):
            self._tools = tools
        elif tools is not None:
            for fn in tools:
                self._tools.register(fn)

        self._lifecycle = AgentState.idle
        self._stop_requested = False

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def tool(self, fn: Callable[..., Any]) -> Callable[..., Any]:
        """Decorator to register a function as a tool on this agent."""
        self._tools.register(fn)
        return fn

    def add_tukuy_tools(self, skills: list[Any]) -> None:
        """Register tukuy skills as tools on this agent.

        Args:
            skills: List of tukuy ``Skill`` instances or ``@skill``-decorated functions.
        """
        self._tools.add_tukuy_skills(skills)

    @property
    def state(self) -> AgentState:
        """Current lifecycle state of the agent."""
        return self._lifecycle

    def stop(self) -> None:
        """Request graceful shutdown after the current iteration."""
        self._stop_requested = True

    @property
    def conversation(self) -> Any:
        """The current persistent conversation, or ``None``."""
        return self._conversation

    @property
    def messages(self) -> list[dict[str, Any]]:
        """Message history from the persistent conversation, or ``[]``."""
        if self._conversation is not None:
            return self._conversation.messages  # type: ignore[no-any-return]
        return []

    def clear_history(self) -> None:
        """Reset the persistent conversation history."""
        if self._conversation is not None:
            self._conversation.clear()

    def as_tool(
        self,
        name: str | None = None,
        description: str | None = None,
        custom_output_extractor: Callable[[AgentResult], str] | None = None,
    ) -> ToolDefinition:
        """Wrap this AsyncAgent as a callable tool for another Agent.

        Creates a :class:`ToolDefinition` whose function accepts a ``prompt``
        string, runs this agent (bridging async to sync), and returns the
        output text.

        Args:
            name: Tool name (defaults to ``self.name`` or ``"agent_tool"``).
            description: Tool description (defaults to ``self.description``).
            custom_output_extractor: Optional function to extract a string
                from :class:`AgentResult`.
        """
        tool_name = name or self.name or "agent_tool"
        tool_desc = description or self.description or f"Run agent {tool_name}"
        agent = self
        extractor = custom_output_extractor

        def _call_agent(prompt: str) -> str:
            """Run the wrapped async agent with the given prompt."""
            try:
                loop = asyncio.get_running_loop()
            except RuntimeError:
                loop = None

            if loop is not None and loop.is_running():
                import concurrent.futures

                with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
                    result = pool.submit(asyncio.run, agent.run(prompt)).result()
            else:
                result = asyncio.run(agent.run(prompt))

            _call_agent._last_agent_result = result  # type: ignore[attr-defined]
            if extractor is not None:
                return extractor(result)
            return result.output_text

        _call_agent._source_agent = agent  # type: ignore[attr-defined]
        _call_agent._last_agent_result = None  # type: ignore[attr-defined]

        return ToolDefinition(
            name=tool_name,
            description=tool_desc,
            parameters={
                "type": "object",
                "properties": {
                    "prompt": {"type": "string", "description": "The prompt to send to the agent"},
                },
                "required": ["prompt"],
            },
            function=_call_agent,
        )

    async def run(
        self,
        prompt: str,
        *,
        deps: Any = None,
        images: list[Any] | None = None,
    ) -> AgentResult:
        """Execute the agent loop to completion (async).

        Creates a fresh conversation, sends the prompt, handles tool calls,
        and optionally parses the final response into ``output_type``.

        Args:
            prompt: The user prompt.
            deps: Optional dependencies.
            images: Optional list of :class:`ImageInput` for vision models.

        Raises:
            RecursionError: If the agent nesting depth exceeds ``max_depth``.
        """
        current_depth = _agent_depth.get()
        if current_depth >= self._max_depth:
            raise RecursionError(f"Agent recursion depth exceeded: {current_depth} >= {self._max_depth}")
        token = _agent_depth.set(current_depth + 1)
        self._lifecycle = AgentState.running
        self._stop_requested = False
        steps: list[AgentStep] = []

        try:
            result = await self._execute(prompt, steps, deps, images=images)
            self._lifecycle = AgentState.idle
            return result
        except Exception:
            self._lifecycle = AgentState.errored
            raise
        finally:
            _agent_depth.reset(token)

    def iter(
        self,
        prompt: str,
        *,
        deps: Any = None,
        images: list[Any] | None = None,
    ) -> AsyncAgentIterator:
        """Execute the agent loop and iterate over steps asynchronously.

        Returns an :class:`AsyncAgentIterator` yielding :class:`AgentStep` objects.
        After iteration, :attr:`AsyncAgentIterator.result` holds the final result.
        """
        gen = self._execute_iter(prompt, deps, images=images)
        return AsyncAgentIterator(gen)

    def run_stream(
        self,
        prompt: str,
        *,
        deps: Any = None,
        images: list[Any] | None = None,
    ) -> AsyncStreamedAgentResult:
        """Execute the agent loop with streaming output (async).

        Returns an :class:`AsyncStreamedAgentResult` yielding :class:`StreamEvent` objects.
        """
        gen = self._execute_stream(prompt, deps, images=images)
        return AsyncStreamedAgentResult(gen)

    # ------------------------------------------------------------------
    # Async callback helpers
    # ------------------------------------------------------------------

    @staticmethod
    async def _invoke_callback(cb: Callable[..., Any], *args: Any) -> Any:
        """Invoke a callback, awaiting it if it's async."""
        if _is_async_callable(cb):
            return await cb(*args)
        return cb(*args)

    # ------------------------------------------------------------------
    # RunContext helpers
    # ------------------------------------------------------------------

    def _build_run_context(
        self,
        prompt: str,
        deps: Any,
        session: UsageSession,
        messages: list[dict[str, Any]],
        iteration: int,
    ) -> RunContext[Any]:
        return RunContext(
            deps=deps,
            model=self._model,
            usage=session.summary(),
            messages=list(messages),
            iteration=iteration,
            prompt=prompt,
        )

    # ------------------------------------------------------------------
    # Tool wrapping (RunContext injection + ModelRetry + callbacks)
    # ------------------------------------------------------------------

    def _wrap_tools_with_context(self, ctx: RunContext[Any], session: UsageSession | None = None) -> ToolRegistry:
        """Return a new :class:`ToolRegistry` with wrapped tool functions.

        All wrappers are **sync** so they work with ``ToolRegistry.execute()``.
        For async tool functions, the wrapper uses
        ``asyncio.get_event_loop().run_until_complete()`` as a fallback.
        If *session* is provided, child agent usage is aggregated into it.
        """
        if not self._tools:
            return ToolRegistry()

        new_registry = ToolRegistry()
        cb = self._agent_callbacks

        for td in self._tools.definitions:
            wants_ctx = _tool_wants_context(td.function)
            original_fn = td.function
            tool_name = td.name
            is_async = _is_async_callable(original_fn)

            def _make_wrapper(
                _fn: Callable[..., Any],
                _wants: bool,
                _name: str,
                _is_async: bool,
                _cb: AgentCallbacks = cb,
                _session: UsageSession | None = session,
            ) -> Callable[..., Any]:
                def _invoke_cb_sync(callback: Callable[..., Any], *cb_args: Any) -> Any:
                    """Invoke a possibly-async callback from a sync tool wrapper."""
                    if _is_async_callable(callback):
                        try:
                            loop = asyncio.get_running_loop()
                        except RuntimeError:
                            loop = None
                        if loop is not None and loop.is_running():
                            import concurrent.futures

                            with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
                                return pool.submit(asyncio.run, callback(*cb_args)).result()
                        else:
                            return asyncio.run(callback(*cb_args))
                    return callback(*cb_args)

                def wrapper(**kwargs: Any) -> Any:
                    if _cb.on_tool_start:
                        _invoke_cb_sync(_cb.on_tool_start, _name, kwargs)
                    try:
                        # Inject skill config via SkillContext for tukuy skills
                        if self._skill_config is not None:
                            _skill_obj = getattr(_fn, "__skill__", None)
                            if _skill_obj is not None:
                                from tukuy import SkillContext

                                kwargs["context"] = SkillContext(config=self._skill_config)

                        # Auto-approve gate: block tukuy skills with side effects
                        if self._auto_approve_safe_only:
                            skill_obj = getattr(_fn, "__skill__", None)
                            if skill_obj is not None:
                                desc = skill_obj.descriptor
                                has_side_effects = getattr(desc, "side_effects", False)
                                has_network = getattr(desc, "requires_network", False)
                                if has_side_effects or has_network:
                                    raise ApprovalRequired(
                                        tool_name=_name,
                                        action="execute tool with side effects",
                                        details={
                                            "side_effects": has_side_effects,
                                            "requires_network": has_network,
                                            "skill_name": desc.name,
                                        },
                                    )

                        if _wants:
                            call_args: tuple[Any, ...] = (ctx,)
                        else:
                            call_args = ()

                        if _is_async:
                            coro = _fn(*call_args, **kwargs)
                            # Try to get running loop; if none, use asyncio.run()
                            try:
                                loop = asyncio.get_running_loop()
                            except RuntimeError:
                                loop = None
                            if loop is not None and loop.is_running():
                                # We're inside an async context — create a new thread
                                import concurrent.futures

                                with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
                                    result = pool.submit(asyncio.run, coro).result()
                            else:
                                result = asyncio.run(coro)
                        else:
                            result = _fn(*call_args, **kwargs)
                    except ApprovalRequired as exc:
                        # Handle approval request
                        if _cb.on_approval_needed:
                            approved = _invoke_cb_sync(_cb.on_approval_needed, exc.tool_name, exc.action, exc.details)
                            if approved:
                                try:
                                    if _is_async:
                                        coro = _fn(*call_args, **kwargs) if not _wants else _fn(ctx, **kwargs)
                                        try:
                                            loop = asyncio.get_running_loop()
                                        except RuntimeError:
                                            loop = None
                                        if loop is not None and loop.is_running():
                                            import concurrent.futures

                                            with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
                                                result = pool.submit(asyncio.run, coro).result()
                                        else:
                                            result = asyncio.run(coro)
                                    else:
                                        if _wants:
                                            result = _fn(ctx, **kwargs)
                                        else:
                                            result = _fn(**kwargs)
                                except ApprovalRequired:
                                    result = f"Error: Tool '{_name}' requires approval but approval was already granted"
                                except ModelRetry as retry_exc:
                                    result = f"Error: {retry_exc.message}"
                            else:
                                result = f"Error: Tool '{_name}' execution denied - approval required: {exc.action}"
                        else:
                            result = f"Error: Tool '{_name}' requires approval but no approval handler is configured"
                    except ModelRetry as exc:
                        result = f"Error: {exc.message}"

                    # Aggregate child agent usage to parent session
                    if _session is not None and hasattr(_fn, "_source_agent"):
                        agent_result = getattr(_fn, "_last_agent_result", None)
                        if agent_result is not None and hasattr(agent_result, "run_usage"):
                            child_usage = agent_result.run_usage
                            child_name = getattr(_fn._source_agent, "name", "") or _name
                            _session.record(
                                {
                                    "meta": {
                                        "prompt_tokens": child_usage.get("prompt_tokens", 0),
                                        "completion_tokens": child_usage.get("completion_tokens", 0),
                                        "total_tokens": child_usage.get("total_tokens", 0),
                                        "cost": child_usage.get("cost", 0.0),
                                    },
                                    "driver": f"sub-agent:{child_name}",
                                }
                            )

                    if _cb.on_tool_end:
                        _invoke_cb_sync(_cb.on_tool_end, _name, result)
                    return result

                return wrapper

            wrapped = _make_wrapper(original_fn, wants_ctx, tool_name, is_async)

            # Build schema: strip RunContext param if present
            params = dict(td.parameters)
            if wants_ctx:
                ctx_param_name = _get_first_param_name(td.function)
                props = dict(params.get("properties", {}))
                props.pop(ctx_param_name, None)
                params = dict(params)
                params["properties"] = props
                req = list(params.get("required", []))
                if ctx_param_name in req:
                    req.remove(ctx_param_name)
                if req:
                    params["required"] = req
                elif "required" in params:
                    del params["required"]

            new_td = ToolDefinition(
                name=td.name,
                description=td.description,
                parameters=params,
                function=wrapped,
            )
            new_registry.add(new_td)

        return new_registry

    # ------------------------------------------------------------------
    # Guardrails
    # ------------------------------------------------------------------

    def _run_input_guardrails(self, ctx: RunContext[Any], prompt: str) -> str:
        for guardrail in self._input_guardrails:
            result = guardrail(ctx, prompt)
            if result is not None:
                prompt = result
        return prompt

    async def _run_output_guardrails(
        self,
        ctx: RunContext[Any],
        result: AgentResult,
        conv: Any,
        session: UsageSession,
        steps: list[AgentStep],
        all_tool_calls: list[dict[str, Any]],
    ) -> AgentResult:
        for guardrail in self._output_guardrails:
            for attempt in range(_OUTPUT_GUARDRAIL_MAX_RETRIES):
                try:
                    guard_result = guardrail(ctx, result)
                    if guard_result is not None:
                        result = guard_result
                    break
                except ModelRetry as exc:
                    if self._is_over_budget(session):
                        break
                    if attempt >= _OUTPUT_GUARDRAIL_MAX_RETRIES - 1:
                        raise ValueError(
                            f"Output guardrail failed after {_OUTPUT_GUARDRAIL_MAX_RETRIES} retries: {exc.message}"
                        ) from exc
                    retry_text = await conv.ask(
                        f"Your response did not pass validation. Error: {exc.message}\n\nPlease try again."
                    )
                    self._extract_steps(
                        conv.messages[-2:],
                        steps,
                        all_tool_calls,
                        getattr(conv, "_full_tool_results", None),
                    )

                    if self._output_type is not None:
                        try:
                            cleaned = clean_json_text(retry_text)
                            parsed = json.loads(cleaned)
                            output = self._output_type.model_validate(parsed)
                        except Exception:
                            output = retry_text
                    else:
                        output = retry_text

                    result = AgentResult(
                        output=output,
                        output_text=retry_text,
                        messages=conv.messages,
                        usage=conv.usage,
                        steps=steps,
                        all_tool_calls=all_tool_calls,
                        state=AgentState.idle,
                        run_usage=session.summary(),
                    )
        return result

    # ------------------------------------------------------------------
    # Budget check
    # ------------------------------------------------------------------

    def _is_over_budget(self, session: UsageSession) -> bool:
        if self._max_cost is not None and session.cost >= self._max_cost:
            return True
        if self._max_tokens is not None and session.total_tokens >= self._max_tokens:
            return True
        return False

    # ------------------------------------------------------------------
    # Internals
    # ------------------------------------------------------------------

    def _resolve_system_prompt(self, ctx: RunContext[Any] | None = None) -> str | None:
        parts: list[str] = []

        if self._system_prompt is not None:
            if isinstance(self._system_prompt, Persona):
                render_kwargs: dict[str, Any] = {}
                if ctx is not None:
                    render_kwargs["model"] = ctx.model
                    render_kwargs["iteration"] = ctx.iteration
                parts.append(self._system_prompt.render(**render_kwargs))
            elif callable(self._system_prompt) and not isinstance(self._system_prompt, str):
                if ctx is not None:
                    parts.append(self._system_prompt(ctx))
                else:
                    parts.append(self._system_prompt(None))
            else:
                parts.append(str(self._system_prompt))

        if self._output_type is not None:
            schema = self._output_type.model_json_schema()
            schema_str = json.dumps(schema, indent=2)
            parts.append(
                "You MUST respond with a single JSON object (no markdown, "
                "no extra text) that validates against this JSON schema:\n"
                f"{schema_str}\n\n"
                "Use double quotes for keys and strings. "
                "If a value is unknown use null."
            )

        return "\n\n".join(parts) if parts else None

    def _build_conversation(
        self,
        system_prompt: str | None = None,
        tools: ToolRegistry | None = None,
        driver_callbacks: DriverCallbacks | None = None,
    ) -> Any:
        """Create or reuse an AsyncConversation for a run.

        When ``persistent_conversation=True`` and a conversation already
        exists, it is reused (with updated tools and callbacks).
        Otherwise a fresh :class:`AsyncConversation` is created.
        """
        from .async_conversation import AsyncConversation

        # Reuse existing conversation in persistent mode
        if self._persistent_conversation and self._conversation is not None:
            conv = self._conversation
            if tools is not None:
                conv._tools = tools
            if driver_callbacks is not None:
                conv._driver.callbacks = driver_callbacks
            return conv

        effective_tools = tools if tools is not None else (self._tools if self._tools else None)

        kwargs: dict[str, Any] = {
            "system_prompt": system_prompt if system_prompt is not None else self._resolve_system_prompt(),
            "tools": effective_tools,
            "max_tool_rounds": self._max_iterations,
        }
        if self._max_tool_result_length is not None:
            kwargs["max_tool_result_length"] = self._max_tool_result_length
        if self._options:
            kwargs["options"] = self._options
        if driver_callbacks is not None:
            kwargs["callbacks"] = driver_callbacks

        if self._driver is not None:
            kwargs["driver"] = self._driver
        else:
            kwargs["model_name"] = self._model

        # Forward budget params when budget enforcement is active
        if self._budget_policy is not None:
            kwargs["budget_policy"] = self._budget_policy
            if self._max_cost is not None:
                kwargs["max_cost"] = self._max_cost
            if self._max_tokens is not None:
                kwargs["max_tokens"] = self._max_tokens
            if self._fallback_models is not None:
                kwargs["fallback_models"] = self._fallback_models
            if self._on_model_fallback is not None:
                kwargs["on_model_fallback"] = self._on_model_fallback

        conv = AsyncConversation(**kwargs)
        if self._persistent_conversation:
            self._conversation = conv
        return conv

    async def _execute(
        self,
        prompt: str,
        steps: list[AgentStep],
        deps: Any,
        *,
        images: list[Any] | None = None,
    ) -> AgentResult:
        """Core async execution: run conversation, extract steps, parse output."""
        from ..infra.tracker import get_tracker

        tracker = get_tracker()

        # 1. Create per-run UsageSession
        session = UsageSession()
        driver_callbacks = DriverCallbacks(
            on_response=session.record,
            on_error=session.record_error,
        )

        # 2. Build initial RunContext
        ctx = self._build_run_context(prompt, deps, session, [], 0)

        # 3. Run input guardrails
        effective_prompt = self._run_input_guardrails(ctx, prompt)

        # 4. Resolve system prompt
        resolved_system_prompt = self._resolve_system_prompt(ctx)

        # 5. Wrap tools with context (pass session for child agent usage aggregation)
        wrapped_tools = self._wrap_tools_with_context(ctx, session)

        # 6. Build AsyncConversation
        conv = self._build_conversation(
            system_prompt=resolved_system_prompt,
            tools=wrapped_tools if wrapped_tools else None,
            driver_callbacks=driver_callbacks,
        )

        # 7. Fire on_iteration callback
        if self._agent_callbacks.on_iteration:
            await self._invoke_callback(self._agent_callbacks.on_iteration, 0)

        # 8. Enforce budget policy at the agent level (pre-call)
        if self._budget_policy is not None:
            state = BudgetState(
                cost_used=session.cost,
                tokens_used=session.total_tokens,
                max_cost=self._max_cost,
                max_tokens=self._max_tokens,
            )
            new_model = enforce_budget(
                state,
                self._budget_policy,
                fallback_models=self._fallback_models,
                current_model=self._model,
                on_model_fallback=self._on_model_fallback,
            )
            if new_model is not None:
                self._model = new_model
                self._driver = None  # force rebuild
                conv = self._build_conversation(
                    system_prompt=resolved_system_prompt,
                    tools=wrapped_tools if wrapped_tools else None,
                    driver_callbacks=driver_callbacks,
                )

        # 9. Ask the conversation (handles full tool loop internally)
        agent_name = self.name or self.__class__.__name__
        with tracker.agent(agent_name):
            t0 = time.perf_counter()
            security_token = None
            if self._security_context is not None:
                from tukuy.safety import set_security_context

                security_token = set_security_context(self._security_context)
            try:
                response_text = await conv.ask(effective_prompt, images=images)
            finally:
                if security_token is not None:
                    from tukuy.safety import reset_security_context

                    reset_security_context(security_token)
            elapsed_ms = (time.perf_counter() - t0) * 1000

            # 9. Extract steps and tool calls
            all_tool_calls: list[dict[str, Any]] = []
            full_results = getattr(conv, "_full_tool_results", None)
            self._extract_steps(conv.messages, steps, all_tool_calls, full_results)

            # Handle output_type parsing
            if self._output_type is not None:
                output, output_text = await self._parse_output(
                    conv, response_text, steps, all_tool_calls, elapsed_ms, session
                )
            else:
                output = response_text
                output_text = response_text

            result = AgentResult(
                output=output,
                output_text=output_text,
                messages=conv.messages,
                usage=conv.usage,
                steps=steps,
                all_tool_calls=all_tool_calls,
                state=AgentState.idle,
                run_usage=session.summary(),
            )

            # 10. Run output guardrails
            if self._output_guardrails:
                result = await self._run_output_guardrails(ctx, result, conv, session, steps, all_tool_calls)

            # 11. Fire callbacks (async-aware)
            if self._agent_callbacks.on_step:
                for step in steps:
                    await self._invoke_callback(self._agent_callbacks.on_step, step)
            if self._agent_callbacks.on_output:
                await self._invoke_callback(self._agent_callbacks.on_output, result)

            if self._agent_callbacks.on_message:
                await self._invoke_callback(self._agent_callbacks.on_message, result.output_text)

            return result

    def _extract_steps(
        self,
        messages: list[dict[str, Any]],
        steps: list[AgentStep],
        all_tool_calls: list[dict[str, Any]],
        full_tool_results: dict[str, str] | None = None,
    ) -> None:
        """Scan conversation messages and populate steps and tool_calls."""
        now = time.time()

        for msg in messages:
            role = msg.get("role", "")
            # Extract usage from message meta if present
            msg_meta = msg.get("meta")
            step_usage = None
            if isinstance(msg_meta, dict):
                step_usage = {
                    "prompt_tokens": msg_meta.get("prompt_tokens", 0),
                    "completion_tokens": msg_meta.get("completion_tokens", 0),
                    "total_tokens": msg_meta.get("total_tokens", 0),
                    "cost": msg_meta.get("cost", 0.0),
                }

            if role == "assistant":
                tc_list = msg.get("tool_calls", [])
                if tc_list:
                    for tc in tc_list:
                        fn = tc.get("function", {})
                        name = fn.get("name", tc.get("name", ""))
                        raw_args = fn.get("arguments", tc.get("arguments", "{}"))
                        if isinstance(raw_args, str):
                            try:
                                args = json.loads(raw_args)
                            except json.JSONDecodeError:
                                args = {}
                        else:
                            args = raw_args

                        steps.append(
                            AgentStep(
                                step_type=StepType.tool_call,
                                timestamp=now,
                                content=msg.get("content", ""),
                                tool_name=name,
                                tool_args=args,
                                usage=step_usage,
                            )
                        )
                        all_tool_calls.append({"name": name, "arguments": args, "id": tc.get("id", "")})
                else:
                    steps.append(
                        AgentStep(
                            step_type=StepType.output,
                            timestamp=now,
                            content=msg.get("content", ""),
                            usage=step_usage,
                        )
                    )

            elif role == "tool":
                tool_call_id = msg.get("tool_call_id")
                # Use the full (pre-truncation) result when available,
                # falling back to the (possibly truncated) message content.
                full_result = None
                if full_tool_results and tool_call_id:
                    full_result = full_tool_results.get(tool_call_id)
                steps.append(
                    AgentStep(
                        step_type=StepType.tool_result,
                        timestamp=now,
                        content=msg.get("content", ""),
                        tool_name=tool_call_id,
                        tool_result=full_result,
                    )
                )

    async def _parse_output(
        self,
        conv: Any,
        response_text: str,
        steps: list[AgentStep],
        all_tool_calls: list[dict[str, Any]],
        elapsed_ms: float,
        session: UsageSession | None = None,
    ) -> tuple[Any, str]:
        """Try to parse ``response_text`` as the output_type, with retries (async)."""
        assert self._output_type is not None

        last_error: Exception | None = None
        text = response_text

        for attempt in range(_OUTPUT_PARSE_MAX_RETRIES):
            try:
                cleaned = clean_json_text(text)
                parsed = json.loads(cleaned)
                model_instance = self._output_type.model_validate(parsed)
                return model_instance, text
            except Exception as exc:
                last_error = exc
                if attempt < _OUTPUT_PARSE_MAX_RETRIES - 1:
                    if session is not None and self._is_over_budget(session):
                        break
                    retry_msg = (
                        f"Your previous response could not be parsed as valid JSON "
                        f"matching the required schema. Error: {exc}\n\n"
                        f"Please try again and respond ONLY with valid JSON."
                    )
                    text = await conv.ask(retry_msg)
                    self._extract_steps(
                        conv.messages[-2:],
                        steps,
                        all_tool_calls,
                        getattr(conv, "_full_tool_results", None),
                    )

        raise ValueError(
            f"Failed to parse output as {self._output_type.__name__} "
            f"after {_OUTPUT_PARSE_MAX_RETRIES} attempts: {last_error}"
        )

    # ------------------------------------------------------------------
    # iter() — async step-by-step
    # ------------------------------------------------------------------

    async def _execute_iter(
        self,
        prompt: str,
        deps: Any,
        *,
        images: list[Any] | None = None,
    ) -> AsyncGenerator[AgentStep, None]:
        """Async generator that executes the agent loop and yields each step."""
        self._lifecycle = AgentState.running
        self._stop_requested = False
        steps: list[AgentStep] = []

        try:
            result = await self._execute(prompt, steps, deps, images=images)
            for step in result.steps:
                yield step
            self._lifecycle = AgentState.idle
            # Store result on the generator for retrieval
            self._last_iter_result = result
        except Exception:
            self._lifecycle = AgentState.errored
            raise

    # ------------------------------------------------------------------
    # run_stream() — async streaming
    # ------------------------------------------------------------------

    async def _execute_stream(
        self,
        prompt: str,
        deps: Any,
        *,
        images: list[Any] | None = None,
    ) -> AsyncGenerator[StreamEvent, None]:
        """Async generator that executes the agent loop and yields stream events.

        Raises:
            RecursionError: If the agent nesting depth exceeds ``max_depth``.
        """
        current_depth = _agent_depth.get()
        if current_depth >= self._max_depth:
            raise RecursionError(f"Agent recursion depth exceeded: {current_depth} >= {self._max_depth}")
        token = _agent_depth.set(current_depth + 1)
        self._lifecycle = AgentState.running
        self._stop_requested = False
        steps: list[AgentStep] = []

        try:
            # 1. Setup
            session = UsageSession()
            driver_callbacks = DriverCallbacks(
                on_response=session.record,
                on_error=session.record_error,
            )
            ctx = self._build_run_context(prompt, deps, session, [], 0)
            effective_prompt = self._run_input_guardrails(ctx, prompt)
            resolved_system_prompt = self._resolve_system_prompt(ctx)
            wrapped_tools = self._wrap_tools_with_context(ctx, session)
            has_tools = bool(wrapped_tools)

            conv = self._build_conversation(
                system_prompt=resolved_system_prompt,
                tools=wrapped_tools if wrapped_tools else None,
                driver_callbacks=driver_callbacks,
            )

            if self._agent_callbacks.on_iteration:
                await self._invoke_callback(self._agent_callbacks.on_iteration, 0)

            security_token = None
            if self._security_context is not None:
                from tukuy.safety import set_security_context

                security_token = set_security_context(self._security_context)
            try:
                if has_tools:
                    response_text = ""
                    async for event in conv.ask_with_tool_events(effective_prompt, images=images):
                        if event["type"] == "tool_call":
                            yield StreamEvent(event_type=StreamEventType.tool_call, data=event)
                        elif event["type"] == "tool_result":
                            yield StreamEvent(event_type=StreamEventType.tool_result, data=event)
                        elif event["type"] == "text_delta":
                            response_text += event["text"]
                            yield StreamEvent(event_type=StreamEventType.text_delta, data=event["text"])
                else:
                    response_text = ""
                    async for chunk in conv.ask_stream(effective_prompt, images=images):
                        response_text += chunk
                        yield StreamEvent(event_type=StreamEventType.text_delta, data=chunk)
            finally:
                if security_token is not None:
                    from tukuy.safety import reset_security_context

                    reset_security_context(security_token)

            # Extract steps
            all_tool_calls: list[dict[str, Any]] = []
            full_results = getattr(conv, "_full_tool_results", None)
            self._extract_steps(conv.messages, steps, all_tool_calls, full_results)

            # Parse output
            if self._output_type is not None:
                output, output_text = await self._parse_output(conv, response_text, steps, all_tool_calls, 0.0, session)
            else:
                output = response_text
                output_text = response_text

            result = AgentResult(
                output=output,
                output_text=output_text,
                messages=conv.messages,
                usage=conv.usage,
                steps=steps,
                all_tool_calls=all_tool_calls,
                state=AgentState.idle,
                run_usage=session.summary(),
            )

            if self._output_guardrails:
                result = await self._run_output_guardrails(ctx, result, conv, session, steps, all_tool_calls)

            if self._agent_callbacks.on_step:
                for step in steps:
                    await self._invoke_callback(self._agent_callbacks.on_step, step)
            if self._agent_callbacks.on_output:
                await self._invoke_callback(self._agent_callbacks.on_output, result)

            if self._agent_callbacks.on_message:
                await self._invoke_callback(self._agent_callbacks.on_message, result.output_text)

            yield StreamEvent(event_type=StreamEventType.output, data=result)

            self._lifecycle = AgentState.idle
            self._last_stream_result = result
        except Exception:
            self._lifecycle = AgentState.errored
            raise
        finally:
            _agent_depth.reset(token)


# ------------------------------------------------------------------
# AsyncAgentIterator
# ------------------------------------------------------------------


class AsyncAgentIterator:
    """Wraps the :meth:`AsyncAgent.iter` async generator, capturing the final result.

    After async iteration completes, :attr:`result` holds the :class:`AgentResult`.
    """

    def __init__(self, gen: AsyncGenerator[AgentStep, None]) -> None:
        self._gen = gen
        self._result: AgentResult | None = None
        self._agent: AsyncAgent[Any] | None = None

    def __aiter__(self) -> AsyncAgentIterator:
        return self

    async def __anext__(self) -> AgentStep:
        try:
            return await self._gen.__anext__()
        except StopAsyncIteration:
            # Try to capture the result from the agent
            ag_frame = getattr(self._gen, "ag_frame", None)
            agent = ag_frame and ag_frame.f_locals.get("self") if ag_frame else None
            if agent and hasattr(agent, "_last_iter_result"):
                self._result = agent._last_iter_result
            raise

    @property
    def result(self) -> AgentResult | None:
        """The final :class:`AgentResult`, available after iteration completes."""
        return self._result


# ------------------------------------------------------------------
# AsyncStreamedAgentResult
# ------------------------------------------------------------------


class AsyncStreamedAgentResult:
    """Wraps the :meth:`AsyncAgent.run_stream` async generator.

    Yields :class:`StreamEvent` objects.  After iteration completes,
    :attr:`result` holds the :class:`AgentResult`.
    """

    def __init__(self, gen: AsyncGenerator[StreamEvent, None]) -> None:
        self._gen = gen
        self._result: AgentResult | None = None

    def __aiter__(self) -> AsyncStreamedAgentResult:
        return self

    async def __anext__(self) -> StreamEvent:
        try:
            event = await self._gen.__anext__()
            # Capture result from the output event
            if event.event_type == StreamEventType.output and isinstance(event.data, AgentResult):
                self._result = event.data
            return event
        except StopAsyncIteration:
            raise

    @property
    def result(self) -> AgentResult | None:
        """The final :class:`AgentResult`, available after iteration completes."""
        return self._result
