"""Strands Agents SDK Python surface.

Wire types live in :mod:`strands.types` (a re-export of the machine-generated
:mod:`strands._generated` package). Records are dataclasses; variant arms are
nested ``VariantCase`` subclasses (e.g. ``StreamEvent.TextDelta``) that pass
``isinstance`` and ``match`` natively.

This module overrides a handful of generated config dataclasses to fold in
ergonomic transforms (seconds → nanoseconds, ``**extras`` → JSON-encoded
``additional_config``, dict → list-of-pairs). Variant-arm wrapping is handled
at the boundary that consumes the value — Agent for its own slots,
SessionManager for storage, McpClientConfig for transport, etc.
"""

from __future__ import annotations

import asyncio
import inspect
import json
import typing
from collections.abc import AsyncIterator, Awaitable, Callable, Iterable
from dataclasses import asdict, is_dataclass
from typing import Any, Protocol, TypeVar, get_type_hints, runtime_checkable

from wasmtime.component import VariantCase as _WitVariantCase

from strands import types

# Hot-path types users import directly from ``strands``. Wire types not
# given an SDK shim below (FileStorage, S3Storage, ...) are passthroughs;
# re-exported here so ``from strands import FileStorage`` works without
# making users reach into ``strands.types``.
from strands.types import (  # noqa: F401
    AgentSkills,
    ContentBlock,
    ContextOffloader,
    CustomStorage,
    FileStorage,
    Interrupt,
    Metrics,
    Role,
    S3Storage,
    SlidingWindowConversationManager,
    StopReason,
    StreamEvent,
    SummarizingConversationManager,
    ToolError,
    Usage,
)


class StrandsError(Exception):
    """Base class for all SDK-raised errors."""


class _ModelErrorBase(StrandsError):
    """Base for errors surfaced by a model provider."""


class ContextWindowOverflowError(_ModelErrorBase):
    """Input exceeded the model's context window and no recovery was possible."""


class MaxTokensError(_ModelErrorBase):
    """Model stopped generating because it hit the max-tokens budget."""

    def __init__(self, message: str, partial_message: types.Message | None = None) -> None:
        super().__init__(message)
        self.partial_message = partial_message


class ModelThrottledError(_ModelErrorBase):
    """Model provider throttled the request. Hook into retry to recover."""


class ProviderTokenCountError(_ModelErrorBase):
    """Provider-native token counting failed; base heuristic should run instead."""


class ToolValidationError(StrandsError):
    """A tool failed validation at registration or invocation time."""


class JsonValidationError(StrandsError):
    """A value could not be serialized to JSON."""


class StructuredOutputError(StrandsError):
    """Model refused to use the structured-output tool even after being forced."""


class ConcurrentInvocationError(StrandsError):
    """Agent is already processing an invocation; concurrent calls are not allowed."""


class SessionError(StrandsError):
    """Session storage read/write failed."""


def _extras_to_json(extras: dict[str, Any] | None) -> str | None:
    return json.dumps(extras) if extras else None


_CONTENT_ARM_BY_TYPE: dict[type, type] = {
    types.TextBlock: types.ContentBlock.Text,
    types.JsonBlock: types.ContentBlock.Json,
    types.ToolUseBlock: types.ContentBlock.ToolUse,
    types.ToolResultBlock: types.ContentBlock.ToolResult,
    types.ReasoningBlock: types.ContentBlock.Reasoning,
    types.CachePointBlock: types.ContentBlock.CachePoint,
    types.ImageBlock: types.ContentBlock.Image,
    types.VideoBlock: types.ContentBlock.Video,
    types.DocumentBlock: types.ContentBlock.Document,
    types.CitationsBlock: types.ContentBlock.Citations,
    types.InterruptResponseBlock: types.ContentBlock.InterruptResponse,
}


def _as_content_block(item: Any) -> Any:
    """Wrap any accepted content shape as a ``ContentBlock`` variant arm."""
    if isinstance(item, str):
        return types.ContentBlock.Text(types.TextBlock(text=item))
    for block_type, arm in _CONTENT_ARM_BY_TYPE.items():
        if isinstance(item, block_type):
            return arm(item)
    return item  # already a ContentBlock variant arm


_MODEL_ARM_BY_TYPE: dict[type, type] = {
    types.BedrockModel: types.ModelConfig.Bedrock,
    types.AnthropicModel: types.ModelConfig.Anthropic,
    types.OpenaiModel: types.ModelConfig.Openai,
    types.GoogleModel: types.ModelConfig.Gemini,
    types.CustomModel: types.ModelConfig.Custom,
}

_CM_ARM_BY_TYPE: dict[type, type] = {
    types.SlidingWindowConversationManager: types.ConversationManagerConfig.SlidingWindow,
    types.SummarizingConversationManager: types.ConversationManagerConfig.Summarizing,
}

_VENDED_TOOL_ARM_BY_TYPE: dict[type, type] = {
    types.BashTool: types.VendedTool.Bash,
    types.FileEditorTool: types.VendedTool.FileEditor,
    types.HttpRequestTool: types.VendedTool.HttpRequest,
    types.NotebookTool: types.VendedTool.Notebook,
}

_VENDED_PLUGIN_ARM_BY_TYPE: dict[type, type] = {
    types.AgentSkills: types.VendedPlugin.Skills,
    types.ContextOffloader: types.VendedPlugin.ContextOffloader,
}


def _wrap(value: Any, arm_table: dict[type, type]) -> Any:
    """Wrap ``value`` in the variant arm whose payload type matches its MRO.

    Walks the MRO so SDK ergonomic subclasses (``BedrockModel`` extends
    ``types.BedrockModel``) hit the same arm as the raw bindgen type. Returns
    the value unchanged if it's already an arm (so passing a fully-constructed
    ``ModelConfig.Bedrock(...)`` is idempotent) or doesn't match anything.
    """
    if value is None or isinstance(value, _WitVariantCase):
        return value
    for cls in type(value).__mro__:
        arm = arm_table.get(cls)
        if arm is not None:
            return arm(value)
    return value


class Message(types.Message):
    def __init__(
        self,
        *,
        role: types.Role,
        content: Iterable[Any],
        metadata: types.MessageMetadata | None = None,
    ) -> None:
        super().__init__(
            role=role,
            content=[_as_content_block(c) for c in content],
            metadata=metadata,
        )

    @classmethod
    def user(cls, *content: Any, metadata: types.MessageMetadata | None = None) -> Message:
        return cls(role=types.Role.USER, content=content, metadata=metadata)

    @classmethod
    def assistant(cls, *content: Any, metadata: types.MessageMetadata | None = None) -> Message:
        return cls(role=types.Role.ASSISTANT, content=content, metadata=metadata)


class BedrockModel(types.BedrockModel):
    def __init__(
        self,
        model_id: str = "us.anthropic.claude-opus-4-7-v1:0",
        *,
        region: str | None = None,
        access_key_id: str | None = None,
        secret_access_key: str | None = None,
        session_token: str | None = None,
        **extras: Any,
    ) -> None:
        # The wasm bundle links the AWS SDK browser build, which has no credential
        # chain. Resolve via botocore so users get the same behavior they'd get
        # from any other Python AWS app (env vars, ~/.aws, SSO, IMDS, etc.).
        if access_key_id is None and secret_access_key is None:
            try:
                import botocore.session

                creds = botocore.session.Session().get_credentials()
                if creds is not None:
                    frozen = creds.get_frozen_credentials()
                    access_key_id = frozen.access_key
                    secret_access_key = frozen.secret_key
                    session_token = frozen.token
            except ImportError:
                pass

        super().__init__(
            model_id=model_id,
            region=region,
            access_key_id=access_key_id,
            secret_access_key=secret_access_key,
            session_token=session_token,
            additional_config=_extras_to_json(extras),
        )


class AnthropicModel(types.AnthropicModel):
    def __init__(self, model_id: str | None = None, *, api_key: str | None = None, **extras: Any) -> None:
        super().__init__(model_id=model_id, api_key=api_key, additional_config=_extras_to_json(extras))


class OpenaiModel(types.OpenaiModel):
    def __init__(self, model_id: str | None = None, *, api_key: str | None = None, **extras: Any) -> None:
        super().__init__(model_id=model_id, api_key=api_key, additional_config=_extras_to_json(extras))


class GoogleModel(types.GoogleModel):
    def __init__(self, model_id: str | None = None, *, api_key: str | None = None, **extras: Any) -> None:
        super().__init__(model_id=model_id, api_key=api_key, additional_config=_extras_to_json(extras))


class CustomModel(types.CustomModel):
    def __init__(
        self,
        provider_id: str,
        *,
        model_id: str | None = None,
        stateful: bool = False,
        **extras: Any,
    ) -> None:
        super().__init__(
            provider_id=provider_id,
            model_id=model_id,
            additional_config=_extras_to_json(extras),
            stateful=stateful,
        )


def _json_default(obj: Any) -> Any:
    if is_dataclass(obj) and not isinstance(obj, type):
        return asdict(obj)
    if hasattr(obj, "__dict__"):
        return {k: v for k, v in obj.__dict__.items() if not k.startswith("_")}
    raise TypeError(f"Object of type {type(obj).__name__} is not JSON serializable")


def agent_node(*, id: str, agent_config: Any, description: str | None, timeout: int | None) -> types.AgentNode:
    """Build an ``AgentNode`` with a JSON-encoded nested agent config.

    ``agent_config`` may be a string (passed through), a dataclass, or any
    object with a ``__dict__``; serialization happens here so call sites
    don't import :mod:`json`.
    """
    encoded = agent_config if isinstance(agent_config, str) else json.dumps(agent_config, default=_json_default)
    return types.AgentNode(id=id, description=description, timeout=timeout, agent_config=encoded)


def multi_agent_node(*, id: str, orchestrator: Any, description: str | None) -> types.MultiAgentNode:
    """Build a ``MultiAgentNode`` with a JSON-encoded nested orchestrator."""
    encoded = orchestrator if isinstance(orchestrator, str) else json.dumps(orchestrator, default=_json_default)
    return types.MultiAgentNode(id=id, description=description, orchestrator=encoded)


class StdioTransport(types.StdioTransport):
    """``StdioTransport`` with ``dict[str, str]`` env shorthand."""

    def __init__(
        self,
        *,
        command: str,
        args: list[str],
        env: dict[str, str],
        cwd: str | None,
    ) -> None:
        super().__init__(
            command=command,
            args=args,
            env=[types.EnvVar(key=k, value=v) for k, v in env.items()],
            cwd=cwd,
        )


class HttpTransport(types.HttpTransport):
    """``HttpTransport`` with ``dict[str, str]`` headers shorthand."""

    def __init__(self, *, url: str, headers: dict[str, str]) -> None:
        super().__init__(
            url=url,
            headers=[types.HttpHeader(name=k, value=v) for k, v in headers.items()],
        )


class SseTransport(types.SseTransport):
    """``SseTransport`` with ``dict[str, str]`` headers shorthand."""

    def __init__(self, *, url: str, headers: dict[str, str]) -> None:
        super().__init__(
            url=url,
            headers=[types.HttpHeader(name=k, value=v) for k, v in headers.items()],
        )


class InterruptResponse(types.RespondArgs):
    """Reply to a paused agent via ``response-stream.respond``."""

    def __init__(self, *, interrupt_id: str, response: Any) -> None:
        payload = response if isinstance(response, str) else json.dumps(response)
        super().__init__(interrupt_id=interrupt_id, response=payload)


class PydanticTool:
    """Tool whose input schema is derived from a pydantic ``BaseModel``."""

    def __init__(
        self,
        *,
        name: str,
        description: str,
        input_model: type,
        func: Callable[..., Any],
    ) -> None:
        if not hasattr(input_model, "model_json_schema") or not hasattr(input_model, "model_validate"):
            raise TypeError(f"input_model must be a pydantic BaseModel subclass; got {input_model!r}")
        self.name = name
        self.description = description
        self._input_model = input_model
        self.input_schema = input_model.model_json_schema()
        self.func = func

    def to_spec(self) -> types.ToolSpec:
        return types.ToolSpec(
            name=self.name,
            description=self.description,
            input_schema=json.dumps(self.input_schema),
        )

    def invoke(self, raw_input: str) -> list[Any]:
        payload = json.loads(raw_input) if raw_input else {}
        validated = self._input_model.model_validate(payload)
        return _coerce_tool_result(self.func(validated))


class Tool:
    """Registered tool: spec plus Python callable."""

    def __init__(
        self,
        *,
        name: str,
        description: str,
        input_schema: dict[str, Any],
        func: Callable[..., Any],
    ) -> None:
        self.name = name
        self.description = description
        self.input_schema = input_schema
        self.func = func

    def to_spec(self) -> types.ToolSpec:
        return types.ToolSpec(
            name=self.name,
            description=self.description,
            input_schema=json.dumps(self.input_schema),
        )

    def invoke(self, raw_input: str) -> list[Any]:
        kwargs = json.loads(raw_input) if raw_input else {}
        return _coerce_tool_result(self.func(**kwargs))


def _coerce_tool_result(result: Any) -> list[Any]:
    if isinstance(result, str):
        return [types.ToolResultContent.Text(types.TextBlock(text=result))]
    if isinstance(result, types.TextBlock):
        return [types.ToolResultContent.Text(result)]
    if isinstance(result, types.JsonBlock):
        return [types.ToolResultContent.Json(result)]
    if isinstance(result, dict):
        return [types.ToolResultContent.Json(types.JsonBlock(json=json.dumps(result)))]
    if is_dataclass(result) and not isinstance(result, type):
        return [types.ToolResultContent.Json(types.JsonBlock(json=json.dumps(asdict(result))))]
    if isinstance(result, list):
        return result
    return [types.ToolResultContent.Text(types.TextBlock(text=str(result)))]


def _py_type_to_schema(py_type: Any) -> dict[str, Any]:
    import types

    origin = typing.get_origin(py_type)

    # Strip Annotated[T, ...] -- only the runtime type matters for the schema.
    if origin is typing.Annotated:
        return _py_type_to_schema(typing.get_args(py_type)[0])

    # Optional[T] / Union[T, None] / T | None: emit T's schema and mark nullable.
    if origin is typing.Union or origin is types.UnionType:
        args = typing.get_args(py_type)
        non_none = [a for a in args if a is not type(None)]
        nullable = len(non_none) != len(args)
        if len(non_none) == 1:
            schema = _py_type_to_schema(non_none[0])
            if nullable:
                schema = {**schema, "nullable": True}
            return schema
        return {}  # heterogeneous union -- caller should supply input_schema

    if py_type is str:
        return {"type": "string"}
    if py_type is int:
        return {"type": "integer"}
    if py_type is float:
        return {"type": "number"}
    if py_type is bool:
        return {"type": "boolean"}
    if origin is list:
        args = typing.get_args(py_type)
        return {"type": "array", "items": _py_type_to_schema(args[0]) if args else {}}
    if origin is dict:
        return {"type": "object"}
    if origin is typing.Literal:
        return {"enum": list(typing.get_args(py_type))}
    return {}


def tool(
    func: Callable[..., Any] | None = None,
    *,
    name: str | None = None,
    description: str | None = None,
) -> Any:
    """Decorator that turns a Python function into a :class:`Tool`."""

    def wrap(f: Callable[..., Any]) -> Tool:
        hints = get_type_hints(f)
        sig = inspect.signature(f)
        properties: dict[str, Any] = {}
        required: list[str] = []
        for param_name, param in sig.parameters.items():
            properties[param_name] = _py_type_to_schema(hints.get(param_name, str))
            if param.default is inspect.Parameter.empty:
                required.append(param_name)
        schema: dict[str, Any] = {"type": "object", "properties": properties}
        if required:
            schema["required"] = required
        return Tool(
            name=name or f.__name__,
            description=description or (f.__doc__ or "").strip() or f.__name__,
            input_schema=schema,
            func=f,
        )

    return wrap(func) if func is not None else wrap


_ToolInput = Tool | PydanticTool | Callable[..., Any]
_ToolChoiceInput = Any  # tagged ToolChoice variant or "name" shorthand or None


def _coerce_tool(item: _ToolInput) -> Tool | PydanticTool:
    if isinstance(item, (Tool, PydanticTool)):
        return item
    if callable(item):
        return tool(item)
    raise TypeError(f"unsupported tool: {type(item).__name__}")


def _coerce_prompt(value: Any) -> Any:
    """Coerce a string or iterable of content blocks to a ``prompt-input``."""
    if isinstance(value, str):
        return value
    if isinstance(value, types.Message):
        raise TypeError(
            "stream_async/invoke take content blocks as input, not Messages. "
            "Pass conversation history via Agent(messages=[...]) instead."
        )
    if hasattr(value, "__iter__") and not isinstance(value, (bytes, str)):
        return [_as_content_block(c) for c in value]
    return value


def _coerce_tool_choice(value: _ToolChoiceInput) -> Any:
    if value is None:
        return None
    if isinstance(value, str):
        return types.ToolChoice.Named(value)
    return value


class Agent:
    """Strands agent. Construct once; call :meth:`invoke` or :meth:`stream_async`."""

    def __init__(
        self,
        *,
        model: types.ModelInput | None = None,
        messages: list[types.Message] | None = None,
        system_prompt: str | list[Any] | None = None,
        tools: list[_ToolInput] | None = None,
        agent_tools: list[types.AgentAsToolConfig] | None = None,
        vended_tools: list[types.VendedToolInput] | None = None,
        vended_plugins: list[types.VendedPluginInput] | None = None,
        mcp_clients: list[types.McpClientConfig] | None = None,
        name: str | None = None,
        id: str | None = None,
        description: str | None = None,
        tool_executor: types.ToolExecutorStrategy | None = None,
        display_output: bool | None = None,
        trace_attributes: list[types.TraceAttribute] | None = None,
        trace_context: types.TraceContext | None = None,
        session: types.SessionManager | None = None,
        conversation_manager: types.ConversationManagerInput | None = None,
        retry: types.RetryConfig | None = None,
        structured_output_schema: str | None = None,
        app_state: dict[str, Any] | None = None,
        model_state: dict[str, Any] | None = None,
    ) -> None:
        self._tools: list[Tool | PydanticTool] = [_coerce_tool(t) for t in (tools or [])]
        identity = None
        if name is not None or id is not None or description is not None:
            identity = types.AgentIdentity(name=name, id=id, description=description)

        wrapped_vended_tools = [_wrap(v, _VENDED_TOOL_ARM_BY_TYPE) for v in vended_tools] if vended_tools else None
        wrapped_vended_plugins = (
            [_wrap(p, _VENDED_PLUGIN_ARM_BY_TYPE) for p in vended_plugins] if vended_plugins else None
        )

        self._config = types.AgentConfig(
            model=_wrap(model, _MODEL_ARM_BY_TYPE),
            model_params=None,
            messages=messages,
            system_prompt=(_coerce_prompt(system_prompt) if system_prompt is not None else None),
            tools=[t.to_spec() for t in self._tools] or None,
            agent_tools=agent_tools,
            vended_tools=wrapped_vended_tools,
            vended_plugins=wrapped_vended_plugins,
            mcp_clients=mcp_clients,
            identity=identity,
            tool_executor=tool_executor,
            display_output=display_output,
            trace_attributes=trace_attributes,
            trace_context=trace_context,
            session=session,
            conversation_manager=_wrap(conversation_manager, _CM_ARM_BY_TYPE),
            retry=retry,
            structured_output_schema=structured_output_schema,
            app_state=json.dumps(app_state) if app_state else None,
            model_state=json.dumps(model_state) if model_state else None,
        )
        self._runtime: Any = None

    @property
    def config(self) -> types.AgentConfig:
        return self._config

    def _ensure_runtime(self) -> Any:
        if self._runtime is None:
            from ._runtime import _AgentRuntime

            self._runtime = _AgentRuntime(self)
        return self._runtime

    async def _ensure_runtime_async(self) -> Any:
        rt = self._ensure_runtime()
        await rt.async_init()
        return rt

    def _lookup_tool(self, name: str) -> Tool | PydanticTool:
        for t in self._tools:
            if getattr(t, "name", None) == name:
                return t
        raise KeyError(f"no tool registered under name {name!r}")

    def _build_invoke_args(
        self,
        prompt: Any,
        tools: list[_ToolInput] | None,
        tool_choice: _ToolChoiceInput,
        structured_output_schema: str | None,
    ) -> types.InvokeArgs:
        extra_tools = [_coerce_tool(t).to_spec() for t in (tools or [])] or None
        return types.InvokeArgs(
            input=_coerce_prompt(prompt),
            tools=extra_tools,
            tool_choice=_coerce_tool_choice(tool_choice),
            structured_output_schema=structured_output_schema,
        )

    async def stream_async(
        self,
        prompt: Any,
        *,
        tools: list[_ToolInput] | None = None,
        tool_choice: _ToolChoiceInput = None,
        structured_output_schema: str | None = None,
    ) -> AsyncIterator[types.StreamEvent]:
        """Yield :class:`StreamEvent` arms as the agent runs."""
        runtime = await self._ensure_runtime_async()
        args = self._build_invoke_args(prompt, tools, tool_choice, structured_output_schema)
        stream = await runtime.generate(args)
        async for event in stream:
            yield event

    async def invoke_async(
        self,
        prompt: Any,
        *,
        tools: list[_ToolInput] | None = None,
        tool_choice: _ToolChoiceInput = None,
        structured_output_schema: str | None = None,
    ) -> AgentResult:
        """Run the agent to completion and return an :class:`AgentResult`."""
        accumulator = _AgentResultAccumulator()
        async for event in self.stream_async(
            prompt,
            tools=tools,
            tool_choice=tool_choice,
            structured_output_schema=structured_output_schema,
        ):
            accumulator.consume(event)
        return accumulator.finalize(self)

    def invoke(
        self,
        prompt: Any,
        *,
        tools: list[_ToolInput] | None = None,
        tool_choice: _ToolChoiceInput = None,
        structured_output_schema: str | None = None,
    ) -> AgentResult:
        """Synchronous wrapper around :meth:`invoke_async`.

        Raises :class:`RuntimeError` if called from a running event loop. Use
        :meth:`invoke_async` directly in Jupyter or async frameworks.
        """
        try:
            asyncio.get_running_loop()
        except RuntimeError:
            pass
        else:
            raise RuntimeError(
                "Agent.invoke() cannot run inside an existing event loop. Use 'await agent.invoke_async(...)' instead."
            )
        return asyncio.run(
            self.invoke_async(
                prompt,
                tools=tools,
                tool_choice=tool_choice,
                structured_output_schema=structured_output_schema,
            )
        )

    def cancel(self) -> None:
        """Cancel the in-flight invocation. Fire-and-forget."""
        if self._runtime is not None:
            self._runtime.cancel()

    async def respond(self, interrupt_id: str, response: Any) -> None:
        runtime = await self._ensure_runtime_async()
        payload = response if isinstance(response, str) else json.dumps(response)
        await runtime.respond(types.RespondArgs(interrupt_id=interrupt_id, response=payload))

    async def get_messages(self) -> list[types.Message]:
        return await (await self._ensure_runtime_async()).get_messages()

    async def set_messages(self, messages: list[types.Message]) -> None:
        await (await self._ensure_runtime_async()).set_messages(messages)

    async def get_app_state(self) -> dict[str, Any]:
        return await (await self._ensure_runtime_async()).get_app_state()

    async def set_app_state(self, state: dict[str, Any]) -> None:
        await (await self._ensure_runtime_async()).set_app_state(state)

    async def get_model_state(self) -> dict[str, Any]:
        return await (await self._ensure_runtime_async()).get_model_state()

    async def set_model_state(self, state: dict[str, Any]) -> None:
        await (await self._ensure_runtime_async()).set_model_state(state)


class _AgentResultAccumulator:
    """Folds the stream of events into the fields of an :class:`AgentResult`."""

    def __init__(self) -> None:
        self._stop: types.StopEvent | None = None
        self._last_message: types.Message | None = None
        self._interrupts: list[types.Interrupt] = []

    def consume(self, event: types.StreamEvent) -> None:
        if isinstance(event, types.StreamEvent.MessageAdded):
            self._last_message = event.value.message
        elif isinstance(event, types.StreamEvent.ModelMessage):
            self._last_message = event.value.message
        elif isinstance(event, types.StreamEvent.Stop):
            self._stop = event.value
        elif isinstance(event, types.StreamEvent.AgentResult):
            self._stop = event.value.stop
        elif isinstance(event, types.StreamEvent.Interrupt):
            self._interrupts.append(event.value)

    def finalize(self, agent: Agent) -> AgentResult:
        stop = self._stop
        last = self._last_message
        if last is None:
            last = types.Message(role=types.Role.ASSISTANT, content=[], metadata=None)
        return AgentResult(
            stop_reason=stop.reason if stop is not None else types.StopReason.END_TURN,
            last_message=last,
            usage=stop.usage if stop is not None else None,
            metrics=None,
            structured_output=(json.loads(stop.structured_output) if stop and stop.structured_output else None),
            interrupts=self._interrupts or None,
        )


_HookEventT = TypeVar("_HookEventT")
_HookCallback = Callable[[Any], Any]


@runtime_checkable
class HookProvider(Protocol):
    """Bundle of related hook registrations."""

    def register_hooks(self, registry: HookRegistry) -> None: ...


class HookRegistry:
    """Register callbacks keyed by ``StreamEvent`` arm class.

    Each arm of the wire ``stream-event`` variant is a distinct Python class
    (``StreamEvent.TextDelta``, ``StreamEvent.Stop``, ...). Subscribers match
    by exact class.

    Callbacks for arms whose name begins with ``After`` dispatch in reverse
    registration order, mirroring teardown semantics. Everything else
    dispatches FIFO.
    """

    def __init__(self) -> None:
        self._callbacks: dict[type, list[_HookCallback]] = {}

    def add_callback(
        self,
        event_type: type[_HookEventT],
        callback: Callable[[_HookEventT], Any],
    ) -> Callable[[], None]:
        entries = self._callbacks.setdefault(event_type, [])
        entry = typing.cast(_HookCallback, callback)
        entries.append(entry)

        def _remove() -> None:
            try:
                self._callbacks[event_type].remove(entry)
            except (KeyError, ValueError):
                pass

        return _remove

    def add_hook(self, provider: HookProvider) -> None:
        provider.register_hooks(self)

    def dispatch(self, event: Any) -> None:
        callbacks = self._callbacks_for(event)
        if any(inspect.iscoroutinefunction(cb) for cb in callbacks):
            raise RuntimeError(f"event={type(event).__name__} | use dispatch_async for async callbacks")
        for cb in callbacks:
            cb(event)

    async def dispatch_async(self, event: Any) -> None:
        for cb in self._callbacks_for(event):
            result = cb(event)
            if inspect.iscoroutine(result):
                await typing.cast(Awaitable[Any], result)

    def _callbacks_for(self, event: Any) -> list[_HookCallback]:
        entries = self._callbacks.get(type(event), [])
        return list(reversed(entries)) if type(event).__name__.startswith("After") else list(entries)


class AgentResult:
    """Terminal result of an agent invocation."""

    def __init__(
        self,
        *,
        stop_reason: types.StopReason,
        last_message: types.Message,
        invocation_state: dict[str, Any] | None = None,
        traces: list[types.AgentTrace] | None = None,
        metrics: types.AgentMetrics | None = None,
        usage: types.Usage | None = None,
        structured_output: Any = None,
        interrupts: list[types.Interrupt] | None = None,
    ) -> None:
        self.stop_reason = stop_reason
        self.last_message = last_message
        self.invocation_state = invocation_state if invocation_state is not None else {}
        self.traces = traces
        self.metrics = metrics
        self.usage = usage
        self.structured_output = structured_output
        self.interrupts = interrupts

    @property
    def context_size(self) -> int | None:
        return self.metrics.latest_context_size if self.metrics else None

    @property
    def projected_context_size(self) -> int | None:
        return self.metrics.projected_context_size if self.metrics else None

    def __str__(self) -> str:
        """Concatenate text from TextBlock and ReasoningBlock content."""
        chunks: list[str] = []
        for block in self.last_message.content:
            tag = getattr(block, "tag", None)
            payload = getattr(block, "payload", None)
            if tag == "text" and payload is not None:
                chunks.append(payload.text)
            elif tag == "reasoning" and payload is not None and payload.text:
                chunks.append(payload.text)
        return "\n".join(chunks)
