from __future__ import annotations

import logging
import warnings
from abc import ABC, abstractmethod
from copy import copy
from typing import (
    TYPE_CHECKING,
    Any,
    Dict,
    List,
    Union,
    Generic,
    TypeVar,
    Callable,
    Iterable,
    Iterator,
    Coroutine,
    AsyncIterator,
    cast,
)
from contextlib import contextmanager, asynccontextmanager
from typing_extensions import Literal, TypedDict, override

import httpx2

from ..._types import Body, Query, Headers, NotGiven
from ..._utils import is_given, consume_sync_iterator, consume_async_iterator
from ...types.beta import (
    BetaMessage,
    BetaMessageParam,
    BetaToolUnionParam,
    BetaRequestToolRemovalBlockParam,
    BetaRequestToolAdditionBlockParam,
)
from ..._base_client import merge_headers
from ._tool_dispatch import tool_registry, tool_error_content, available_tool_names
from ._beta_functions import (
    ToolError,
    BetaFunctionTool,
    BetaRunnableTool,
    BetaAsyncFunctionTool,
    BetaAsyncRunnableTool,
    BetaBuiltinFunctionTool,
    BetaAsyncBuiltinFunctionTool,
)
from .._stainless_helpers import stainless_helper_header
from ..streaming._beta_messages import BetaMessageStream, BetaAsyncMessageStream
from ...types.beta.beta_stop_reason import BetaStopReason
from ...types.beta.parsed_beta_message import ResponseFormatT, ParsedBetaMessage, ParsedBetaContentBlock
from ...types.beta.message_create_params import ParseMessageCreateParamsBase
from ...types.beta.beta_output_config_param import BetaOutputConfigParam
from ...types.beta.beta_compaction_config_param import BetaCompactionConfigParam
from ...types.beta.beta_tool_result_block_param import BetaToolResultBlockParam

if TYPE_CHECKING:
    from ..._client import Anthropic, AsyncAnthropic


AnyFunctionToolT = TypeVar(
    "AnyFunctionToolT",
    bound=Union[
        BetaFunctionTool[Any], BetaAsyncFunctionTool[Any], BetaBuiltinFunctionTool, BetaAsyncBuiltinFunctionTool
    ],
)
RunnerItemT = TypeVar("RunnerItemT")


log = logging.getLogger(__name__)

_Step = Literal["run_tools", "resume", "stop"]

# Every stop reason maps to exactly one step. The runner tests assert this mapping
# covers `BetaStopReason`, so a newly generated value must be classified here.
_STOP_REASON_STEPS: dict[BetaStopReason, _Step] = {
    "tool_use": "run_tools",
    "pause_turn": "resume",
    # pause_after_compaction hands the turn back before the model answers; sending it back unchanged continues it.
    "compaction": "resume",
    "end_turn": "stop",
    "stop_sequence": "stop",
    "max_tokens": "stop",
    "model_context_window_exceeded": "stop",
    "refusal": "stop",
}


def _determine_next_step_from_stop_reason(stop_reason: BetaStopReason | None) -> _Step:
    """Decide how the runner loop treats a finished assistant turn.

    - `run_tools`: run the turn's client tool calls, append their results and continue; stop if there are none.
    - `resume`: the turn is not finished; send it back unchanged, running no tool calls, so the server continues it.
    - `stop`: terminal; the turn is the final message and its tool_use blocks must not be executed.
    """
    if stop_reason is not None and stop_reason in _STOP_REASON_STEPS:
        return _STOP_REASON_STEPS[stop_reason]
    # Absent and unknown (forward-compatible) values stop like any other finished turn.
    return "stop"


def _reject_compaction_param(params: ParseMessageCreateParamsBase[Any]) -> None:
    compaction = params.get("compaction")
    if compaction is not None and is_given(compaction):
        raise ValueError(
            "`compaction` cannot be set on a tool runner: every request in the loop would compact again. "
            "Call `runner.compact_before_next_turn()` when the conversation should be compacted instead."
        )


def _without_format(output_config: BetaOutputConfigParam) -> Dict[str, Any]:
    return {key: value for key, value in output_config.items() if key != "format"}


def _without_compaction_incompatible_params(params: ParseMessageCreateParamsBase[Any]) -> Dict[str, Any]:
    """A compaction request returns only the compaction block, never a reply, so the API rejects the
    params that only shape a reply. The runner's later requests keep them.
    """
    trimmed: Dict[str, Any] = {**params}
    for name in ("context_management", "stop_sequences", "output_format"):
        trimmed.pop(name, None)
    tool_choice = params.get("tool_choice")
    if is_given(tool_choice) and tool_choice and tool_choice["type"] in ("any", "tool"):
        del trimmed["tool_choice"]
    output_config = params.get("output_config")
    if is_given(output_config) and output_config:
        trimmed["output_config"] = _without_format(output_config)
    fallbacks = params.get("fallbacks")
    if is_given(fallbacks) and fallbacks and not isinstance(fallbacks, str):
        trimmed["fallbacks"] = [
            {**fallback, "output_config": _without_format(fallback_output_config)}
            if (fallback_output_config := fallback.get("output_config"))
            else fallback
            for fallback in fallbacks
        ]
    return trimmed


class RequestOptions(TypedDict, total=False):
    extra_headers: Headers | None
    extra_query: Query | None
    extra_body: Body | None
    timeout: float | httpx2.Timeout | None | NotGiven


class BaseToolRunner(Generic[AnyFunctionToolT, ResponseFormatT]):
    def __init__(
        self,
        *,
        params: ParseMessageCreateParamsBase[ResponseFormatT],
        options: RequestOptions,
        tools: Iterable[AnyFunctionToolT],
        max_iterations: int | None = None,
    ) -> None:
        _reject_compaction_param(params)
        self._tools_by_name = tool_registry(tools)
        self._params: ParseMessageCreateParamsBase[ResponseFormatT] = {
            **params,
            "messages": [message for message in params["messages"]],
        }
        helper_header = stainless_helper_header(
            tools=self._tools_by_name.values(),
            messages=params.get("messages"),
        )
        if helper_header:
            merged_headers = merge_headers(helper_header, options.get("extra_headers") or {})
            options = {**options, "extra_headers": merged_headers}
        self._options = options
        self._messages_modified = False
        self._cached_tool_call_response: BetaMessageParam | None = None
        self._max_iterations = max_iterations
        self._iteration_count = 0
        self._pending_compaction: BetaCompactionConfigParam | None = None
        self._messages_being_compacted: Iterable[BetaMessageParam] | None = None
        self._pending_tool_changes: list[BetaRequestToolAdditionBlockParam | BetaRequestToolRemovalBlockParam] = []

    def set_messages_params(
        self,
        params: ParseMessageCreateParamsBase[ResponseFormatT]
        | Callable[[ParseMessageCreateParamsBase[ResponseFormatT]], ParseMessageCreateParamsBase[ResponseFormatT]],
    ) -> None:
        """
        Update the parameters for the next API call. This invalidates any cached tool responses.

        Args:
            params (ParsedMessageCreateParamsBase[ResponseFormatT] | Callable): Either new parameters or a function to mutate existing parameters
        """
        if callable(params):
            params = params(self._params)
        _reject_compaction_param(params)
        if self._messages_being_compacted is not None and params["messages"] is not self._messages_being_compacted:
            raise ValueError(
                "Message params can't be changed while the conversation is being compacted, because the compaction "
                "response replaces them. Make the change on the next iteration."
            )
        if self._pending_compaction is not None or self._messages_being_compacted is not None:
            self._check_can_compact(params)
        self._params = params

    def append_messages(self, *messages: BetaMessageParam | ParsedBetaMessage[ResponseFormatT]) -> None:
        """Add one or more messages to the conversation history.

        This invalidates the cached tool response, i.e. if tools were already called, then they will
        be called again on the next loop iteration.
        """
        message_params: List[BetaMessageParam] = [
            message.to_param() if isinstance(message, BetaMessage) else message for message in messages
        ]
        self.set_messages_params(lambda params: {**params, "messages": [*params["messages"], *message_params]})
        self._messages_modified = True
        self._cached_tool_call_response = None

    def compact_before_next_turn(self, compaction: BetaCompactionConfigParam | None = None) -> None:
        """Compact the conversation before the model's next turn.

        This only schedules the compaction. Once the current turn has finished, including
        any tool calls, the runner requests a summary and replaces the message history with
        the compaction response the API returns. That response is yielded like any other
        message, with `stop_reason == "compaction"`, and the runner then carries on.

        Calling this again before the compaction runs replaces the pending one. Requires
        the `compact-2026-09-04` beta.

        Args:
            compaction: The compaction config to send. Defaults to `{"type": "summarize"}`.
        """
        if self._messages_being_compacted is not None:
            return
        self._check_can_compact(self._params)
        self._pending_compaction = {"type": "summarize"} if compaction is None else compaction

    def _check_can_compact(self, params: ParseMessageCreateParamsBase[ResponseFormatT]) -> None:
        # The compaction request is sent without `context_management`, so the API can't reject this
        # combination there: it would run and bill the compaction, then reject the next request, where
        # the compaction response and the compaction edit meet.
        context_management = params.get("context_management")
        edits = context_management.get("edits") if is_given(context_management) and context_management else None
        if any(edit["type"].startswith("compact_") for edit in edits or ()):
            raise ValueError(
                "`compact_before_next_turn()` can't be used while `context_management` has a compaction edit, "
                "because the API doesn't accept a compaction block together with one. Remove the edit first."
            )

    def _pop_compaction_request_params(
        self, compaction: BetaCompactionConfigParam
    ) -> ParseMessageCreateParamsBase[ResponseFormatT]:
        self._check_can_compact(self._params)
        params = {**_without_compaction_incompatible_params(self._params), "compaction": compaction}
        self._pending_compaction = None
        self._messages_being_compacted = self._params["messages"]
        return cast("ParseMessageCreateParamsBase[ResponseFormatT]", params)

    def _prepare_compaction_after_final_turn(
        self, message: ParsedBetaMessage[ResponseFormatT]
    ) -> BetaCompactionConfigParam | None:
        compaction = self._pending_compaction
        if compaction is not None and not self._messages_modified:
            if any(block.type == "tool_use" for block in message.content):
                # A turn that was cut short can end with tool calls that are never run, and the API
                # can't compact a conversation whose last turn has an unanswered tool call.
                log.warning(
                    "The pending compaction was skipped because the last turn (stop_reason=%r) ended with tool calls "
                    "that were not run. Call `compact_before_next_turn()` again if you continue the conversation.",
                    message.stop_reason,
                )
                self._pending_compaction = None
                return None
            self.append_messages(message)
        return compaction

    def _register_compaction_response(self, message: ParsedBetaMessage[ResponseFormatT]) -> bool:
        # Summary or not, the request went out with any edit made before it.
        self._messages_modified = False
        if not any(block.type == "compaction" and block.content for block in message.content):
            log.warning("Compaction produced no summary; keeping the conversation as it is.")
            return False
        for name in self._tools_by_name.keys() - self._available_tool_names():
            del self._tools_by_name[name]
        # The response has to be sent back as it came, first, replacing the messages it summarizes.
        self._params = {**self._params, "messages": [message.to_param()]}
        return True

    def add_tools(self, *tools: AnyFunctionToolT | BetaToolUnionParam) -> None:
        """Give the model more tools without changing the `tools` param, which would miss the prompt cache.

        The definitions are sent in `tool_addition` blocks with the next request. A function tool is run straight
        away, in place of any tool of the same name, even for a call already in the message being handled. A raw
        definition is for server tools, such as web search: the tool runner never runs it, and it stops running a
        function tool of the same name. Requires the `inline-tools-2026-09-15` beta.

        Args:
            *tools: Function tools, such as `@beta_tool` functions, or raw tool definitions.
        """
        for tool in tools:
            if isinstance(tool, dict):
                definition: BetaToolUnionParam = copy(tool)
                name = tool.get("name")
                if isinstance(name, str):
                    self._tools_by_name.pop(name, None)
            else:
                definition = tool.to_dict()
                self._tools_by_name[tool.name] = tool
            self._pending_tool_changes.append(
                {"type": "tool_addition", "tool": {"type": "tool_definition", "definition": definition}}
            )

    def remove_tools(self, *tools: AnyFunctionToolT | str) -> None:
        """Take tools away from the model without changing the `tools` param, which would miss the prompt cache.

        The tools stop being run at once, and the model is told in `tool_removal` blocks with the next request.
        Requires the `inline-tools-2026-09-15` beta.

        Args:
            *tools: The tools to remove, or their names.
        """
        for tool in tools:
            name = tool if isinstance(tool, str) else tool.name
            self._tools_by_name.pop(name, None)
            self._pending_tool_changes.append(
                {"type": "tool_removal", "tool": {"type": "tool_reference", "name": name}}
            )

    def _send_pending_tool_changes(self, hold: bool) -> None:
        # A turn that stopped on `pause_turn` is sent back to be continued, so it has to stay last.
        if hold or not self._pending_tool_changes:
            return
        # Not `append_messages()`: that would make the runner leave this turn's messages for the caller to append.
        self._params = {
            **self._params,
            "messages": [*self._params["messages"], {"role": "system", "content": self._pending_tool_changes}],
        }
        self._pending_tool_changes = []

    def _should_stop(self) -> bool:
        if self._max_iterations is not None and self._iteration_count >= self._max_iterations:
            return True
        return False

    def _available_tool_names(self) -> set[str]:
        """The tool names currently available, after applying any
        mid-conversation `tool_removal` / `tool_addition` blocks.

        Removal is only a hint to the model, which can still emit a `tool_use`
        for a withdrawn tool; a name absent from this set routes that call down
        the same unknown-tool path as a tool that was never declared.
        """
        # Changes made since the last request are not in the history yet.
        pending: BetaMessageParam = {"role": "system", "content": self._pending_tool_changes}
        return available_tool_names([*self._params["messages"], pending], self._tools_by_name)


class BaseSyncToolRunner(BaseToolRunner[BetaRunnableTool, ResponseFormatT], Generic[RunnerItemT, ResponseFormatT], ABC):
    def __init__(
        self,
        *,
        params: ParseMessageCreateParamsBase[ResponseFormatT],
        options: RequestOptions,
        tools: Iterable[BetaRunnableTool],
        client: Anthropic,
        max_iterations: int | None = None,
    ) -> None:
        super().__init__(
            params=params,
            options=options,
            tools=tools,
            max_iterations=max_iterations,
        )
        self._client = client
        self._iterator = self.__run__()
        self._last_message: (
            Callable[[], ParsedBetaMessage[ResponseFormatT]] | ParsedBetaMessage[ResponseFormatT] | None
        ) = None

    def __next__(self) -> RunnerItemT:
        return self._iterator.__next__()

    def __iter__(self) -> Iterator[RunnerItemT]:
        for item in self._iterator:
            yield item

    @abstractmethod
    @contextmanager
    def _handle_request(self, params: ParseMessageCreateParamsBase[ResponseFormatT]) -> Iterator[RunnerItemT]:
        raise NotImplementedError()
        yield  # type: ignore[unreachable]

    def _compact(self, compaction: BetaCompactionConfigParam) -> Iterator[RunnerItemT]:
        last_message = self._last_message
        try:
            with self._handle_request(self._pop_compaction_request_params(compaction)) as item:
                yield item
                message = self._get_last_message()
                assert message is not None
        finally:
            self._messages_being_compacted = None
        if not self._register_compaction_response(message):
            # `until_done()` still returns the turn the run ended on.
            self._last_message = last_message

    def _compact_after_final_turn(self, message: ParsedBetaMessage[ResponseFormatT]) -> Iterator[RunnerItemT]:
        compaction = self._prepare_compaction_after_final_turn(message)
        if compaction is not None:
            yield from self._compact(compaction)

    def __run__(self) -> Iterator[RunnerItemT]:
        stop_reason: BetaStopReason | None = None
        while not self._should_stop():
            self._send_pending_tool_changes(hold=stop_reason == "pause_turn")
            # The API can't compact a conversation that ends mid-turn, so a paused turn is resumed first.
            turn_paused = _determine_next_step_from_stop_reason(stop_reason) == "resume"
            compaction = None if turn_paused else self._pending_compaction
            if compaction is not None:
                yield from self._compact(compaction)
                continue

            with self._handle_request(self._params) as item:
                yield item
                message = self._get_last_message()
                assert message is not None

                # Update container from response for programmatic tool calling support
                last_assistant_message = self._get_last_assistant_message()
                if last_assistant_message is not None and last_assistant_message.container is not None:
                    self._params["container"] = last_assistant_message.container.id

            self._iteration_count += 1

            stop_reason = message.stop_reason
            next_step = _determine_next_step_from_stop_reason(stop_reason)
            if next_step == "stop":
                log.debug("Turn ended with stop_reason %r, exiting from tool runner loop.", message.stop_reason)
                yield from self._compact_after_final_turn(message)
                return

            if next_step == "resume":
                if not self._messages_modified:
                    self.append_messages(message)
            else:
                response = self.generate_tool_call_response()
                if response is None:
                    log.debug("Tool call was not requested, exiting from tool runner loop.")
                    yield from self._compact_after_final_turn(message)
                    return
                if not self._messages_modified:
                    self.append_messages(message, response)

            self._messages_modified = False
            self._cached_tool_call_response = None

    def until_done(self) -> ParsedBetaMessage[ResponseFormatT]:
        """
        Consumes the tool runner stream and returns the last message if it has not been consumed yet.
        If it has, it simply returns the last message.
        """
        consume_sync_iterator(self)
        last_message = self._get_last_message()
        assert last_message is not None
        return last_message

    def generate_tool_call_response(self) -> BetaMessageParam | None:
        """Generate a MessageParam by calling tool functions with any tool use blocks from the last message.

        Note the tool call response is cached, repeated calls to this method will return the same response.

        None can be returned if no tool call was applicable.
        """
        if self._cached_tool_call_response is not None:
            log.debug("Returning cached tool call response.")
            return self._cached_tool_call_response
        response = self._generate_tool_call_response()
        self._cached_tool_call_response = response
        return response

    def _generate_tool_call_response(self) -> BetaMessageParam | None:
        content = self._get_last_assistant_message_content()
        if not content:
            return None

        tool_use_blocks = [block for block in content if block.type == "tool_use"]
        if not tool_use_blocks:
            return None

        results: list[BetaToolResultBlockParam] = []
        available = self._available_tool_names()

        for tool_use in tool_use_blocks:
            tool = self._tools_by_name.get(tool_use.name) if tool_use.name in available else None
            if tool is None:
                warnings.warn(
                    f"Tool '{tool_use.name}' not found in tool runner. "
                    f"Available tools: {list(self._tools_by_name.keys())}. "
                    f"If using a raw tool definition, handle the tool call manually and use `append_messages()` to add the result. "
                    f"Otherwise, pass the tool using `beta_tool(func)` or a `@beta_tool` decorated function.",
                    UserWarning,
                    stacklevel=3,
                )
                results.append(
                    {
                        "type": "tool_result",
                        "tool_use_id": tool_use.id,
                        "content": f"Error: Tool '{tool_use.name}' not found",
                        "is_error": True,
                    }
                )
                continue

            try:
                result = tool.call(tool_use.input)
                results.append({"type": "tool_result", "tool_use_id": tool_use.id, "content": result})
            except ToolError as exc:
                results.append(
                    {
                        "type": "tool_result",
                        "tool_use_id": tool_use.id,
                        "content": tool_error_content(exc),
                        "is_error": True,
                    }
                )
            except Exception as exc:
                log.exception(f"Error occurred while calling tool: {tool.name}", exc_info=exc)
                results.append(
                    {
                        "type": "tool_result",
                        "tool_use_id": tool_use.id,
                        "content": tool_error_content(exc),
                        "is_error": True,
                    }
                )

        return {"role": "user", "content": results}

    def _get_last_message(self) -> ParsedBetaMessage[ResponseFormatT] | None:
        if callable(self._last_message):
            return self._last_message()
        return self._last_message

    def _get_last_assistant_message(self) -> ParsedBetaMessage[ResponseFormatT] | None:
        last_message = self._get_last_message()
        if last_message is None or last_message.role != "assistant" or not last_message.content:
            return None

        return last_message

    def _get_last_assistant_message_content(self) -> list[ParsedBetaContentBlock[ResponseFormatT]] | None:
        last_assistant_message = self._get_last_assistant_message()
        if last_assistant_message is None:
            return None

        return last_assistant_message.content


class BetaToolRunner(BaseSyncToolRunner[ParsedBetaMessage[ResponseFormatT], ResponseFormatT]):
    @override
    @contextmanager
    def _handle_request(
        self, params: ParseMessageCreateParamsBase[ResponseFormatT]
    ) -> Iterator[ParsedBetaMessage[ResponseFormatT]]:
        message = self._client.beta.messages.parse(**params, **self._options)
        self._last_message = message
        yield message


class BetaStreamingToolRunner(BaseSyncToolRunner[BetaMessageStream[ResponseFormatT], ResponseFormatT]):
    @override
    @contextmanager
    def _handle_request(
        self, params: ParseMessageCreateParamsBase[ResponseFormatT]
    ) -> Iterator[BetaMessageStream[ResponseFormatT]]:
        with self._client.beta.messages.stream(**params, **self._options) as stream:
            self._last_message = stream.get_final_message
            yield stream


class BaseAsyncToolRunner(
    BaseToolRunner[BetaAsyncRunnableTool, ResponseFormatT], Generic[RunnerItemT, ResponseFormatT], ABC
):
    def __init__(
        self,
        *,
        params: ParseMessageCreateParamsBase[ResponseFormatT],
        options: RequestOptions,
        tools: Iterable[BetaAsyncRunnableTool],
        client: AsyncAnthropic,
        max_iterations: int | None = None,
    ) -> None:
        super().__init__(
            params=params,
            options=options,
            tools=tools,
            max_iterations=max_iterations,
        )
        self._client = client
        self._iterator = self.__run__()
        self._last_message: (
            Callable[[], Coroutine[None, None, ParsedBetaMessage[ResponseFormatT]]]
            | ParsedBetaMessage[ResponseFormatT]
            | None
        ) = None

    async def __anext__(self) -> RunnerItemT:
        return await self._iterator.__anext__()

    async def __aiter__(self) -> AsyncIterator[RunnerItemT]:
        async for item in self._iterator:
            yield item

    @abstractmethod
    @asynccontextmanager
    async def _handle_request(
        self, params: ParseMessageCreateParamsBase[ResponseFormatT]
    ) -> AsyncIterator[RunnerItemT]:
        raise NotImplementedError()
        yield  # type: ignore[unreachable]

    async def _compact(self, compaction: BetaCompactionConfigParam) -> AsyncIterator[RunnerItemT]:
        last_message = self._last_message
        try:
            async with self._handle_request(self._pop_compaction_request_params(compaction)) as item:
                yield item
                message = await self._get_last_message()
                assert message is not None
        finally:
            self._messages_being_compacted = None
        if not self._register_compaction_response(message):
            # `until_done()` still returns the turn the run ended on.
            self._last_message = last_message

    async def _compact_after_final_turn(
        self, message: ParsedBetaMessage[ResponseFormatT]
    ) -> AsyncIterator[RunnerItemT]:
        compaction = self._prepare_compaction_after_final_turn(message)
        if compaction is not None:
            async for item in self._compact(compaction):
                yield item

    async def __run__(self) -> AsyncIterator[RunnerItemT]:
        stop_reason: BetaStopReason | None = None
        while not self._should_stop():
            self._send_pending_tool_changes(hold=stop_reason == "pause_turn")
            # The API can't compact a conversation that ends mid-turn, so a paused turn is resumed first.
            turn_paused = _determine_next_step_from_stop_reason(stop_reason) == "resume"
            compaction = None if turn_paused else self._pending_compaction
            if compaction is not None:
                async for item in self._compact(compaction):
                    yield item
                continue

            async with self._handle_request(self._params) as item:
                yield item
                message = await self._get_last_message()
                assert message is not None

                # Update container from response for programmatic tool calling support
                last_assistant_message = await self._get_last_assistant_message()
                if last_assistant_message is not None and last_assistant_message.container is not None:
                    self._params["container"] = last_assistant_message.container.id

            self._iteration_count += 1

            stop_reason = message.stop_reason
            next_step = _determine_next_step_from_stop_reason(stop_reason)
            if next_step == "stop":
                log.debug("Turn ended with stop_reason %r, exiting from tool runner loop.", message.stop_reason)
                async for item in self._compact_after_final_turn(message):
                    yield item
                return

            if next_step == "resume":
                if not self._messages_modified:
                    self.append_messages(message)
            else:
                response = await self.generate_tool_call_response()
                if response is None:
                    log.debug("Tool call was not requested, exiting from tool runner loop.")
                    async for item in self._compact_after_final_turn(message):
                        yield item
                    return
                if not self._messages_modified:
                    self.append_messages(message, response)

            self._messages_modified = False
            self._cached_tool_call_response = None

    async def until_done(self) -> ParsedBetaMessage[ResponseFormatT]:
        """
        Consumes the tool runner stream and returns the last message if it has not been consumed yet.
        If it has, it simply returns the last message.
        """
        await consume_async_iterator(self)
        last_message = await self._get_last_message()
        assert last_message is not None
        return last_message

    async def generate_tool_call_response(self) -> BetaMessageParam | None:
        """Generate a MessageParam by calling tool functions with any tool use blocks from the last message.

        Note the tool call response is cached, repeated calls to this method will return the same response.

        None can be returned if no tool call was applicable.
        """
        if self._cached_tool_call_response is not None:
            log.debug("Returning cached tool call response.")
            return self._cached_tool_call_response

        response = await self._generate_tool_call_response()
        self._cached_tool_call_response = response
        return response

    async def _get_last_message(self) -> ParsedBetaMessage[ResponseFormatT] | None:
        if callable(self._last_message):
            return await self._last_message()
        return self._last_message

    async def _get_last_assistant_message(self) -> ParsedBetaMessage[ResponseFormatT] | None:
        last_message = await self._get_last_message()
        if last_message is None or last_message.role != "assistant" or not last_message.content:
            return None

        return last_message

    async def _get_last_assistant_message_content(self) -> list[ParsedBetaContentBlock[ResponseFormatT]] | None:
        last_assistant_message = await self._get_last_assistant_message()
        if last_assistant_message is None:
            return None

        return last_assistant_message.content

    async def _generate_tool_call_response(self) -> BetaMessageParam | None:
        content = await self._get_last_assistant_message_content()
        if not content:
            return None

        tool_use_blocks = [block for block in content if block.type == "tool_use"]
        if not tool_use_blocks:
            return None

        results: list[BetaToolResultBlockParam] = []
        available = self._available_tool_names()

        for tool_use in tool_use_blocks:
            tool = self._tools_by_name.get(tool_use.name) if tool_use.name in available else None
            if tool is None:
                warnings.warn(
                    f"Tool '{tool_use.name}' not found in tool runner. "
                    f"Available tools: {list(self._tools_by_name.keys())}. "
                    f"If using a raw tool definition, handle the tool call manually and use `append_messages()` to add the result. "
                    f"Otherwise, pass the tool using `beta_async_tool(func)` or a `@beta_async_tool` decorated function.",
                    UserWarning,
                    stacklevel=3,
                )
                results.append(
                    {
                        "type": "tool_result",
                        "tool_use_id": tool_use.id,
                        "content": f"Error: Tool '{tool_use.name}' not found",
                        "is_error": True,
                    }
                )
                continue

            try:
                result = await tool.call(tool_use.input)
                results.append({"type": "tool_result", "tool_use_id": tool_use.id, "content": result})
            except ToolError as exc:
                results.append(
                    {
                        "type": "tool_result",
                        "tool_use_id": tool_use.id,
                        "content": tool_error_content(exc),
                        "is_error": True,
                    }
                )
            except Exception as exc:
                log.exception(f"Error occurred while calling tool: {tool.name}", exc_info=exc)
                results.append(
                    {
                        "type": "tool_result",
                        "tool_use_id": tool_use.id,
                        "content": tool_error_content(exc),
                        "is_error": True,
                    }
                )

        return {"role": "user", "content": results}


class BetaAsyncToolRunner(BaseAsyncToolRunner[ParsedBetaMessage[ResponseFormatT], ResponseFormatT]):
    @override
    @asynccontextmanager
    async def _handle_request(
        self, params: ParseMessageCreateParamsBase[ResponseFormatT]
    ) -> AsyncIterator[ParsedBetaMessage[ResponseFormatT]]:
        message = await self._client.beta.messages.parse(**params, **self._options)
        self._last_message = message
        yield message


class BetaAsyncStreamingToolRunner(BaseAsyncToolRunner[BetaAsyncMessageStream[ResponseFormatT], ResponseFormatT]):
    @override
    @asynccontextmanager
    async def _handle_request(
        self, params: ParseMessageCreateParamsBase[ResponseFormatT]
    ) -> AsyncIterator[BetaAsyncMessageStream[ResponseFormatT]]:
        async with self._client.beta.messages.stream(**params, **self._options) as stream:
            self._last_message = stream.get_final_message
            yield stream