"""Tool middleware chain."""
from __future__ import annotations

from typing import Callable, Awaitable

from logos_adk.types import ToolCall, ToolResult

ToolExecutor = Callable[[ToolCall], Awaitable[ToolResult]]
ToolMiddleware = Callable[[ToolCall, ToolExecutor], Awaitable[ToolResult]]


class ToolMiddlewareChain:
    """Chain of middleware wrapping tool execution."""

    def __init__(self, executor: ToolExecutor) -> None:
        self._executor = executor
        self._middlewares: list[ToolMiddleware] = []

    def use(self, middleware: ToolMiddleware) -> None:
        """Add a middleware. First added = outermost (executes first)."""
        self._middlewares.append(middleware)

    async def execute(self, call: ToolCall) -> ToolResult:
        """Execute the tool call through the middleware chain."""
        chain = self._executor
        for mw in reversed(self._middlewares):
            chain = _wrap(mw, chain)
        return await chain(call)


def _wrap(mw: ToolMiddleware, next_fn: ToolExecutor) -> ToolExecutor:
    async def wrapped(call: ToolCall) -> ToolResult:
        return await mw(call, next_fn)
    return wrapped
