"""Tool registry — register, list, lookup, and execute tools."""

from __future__ import annotations

from typing import TYPE_CHECKING, Any

from koan.tools.executor import execute_tool

if TYPE_CHECKING:
    from collections.abc import Callable

    from koan.tools.types import ToolDefinition, ToolResult


class ToolRegistry:
    """Registry for tool functions.

    Usage::

        registry = ToolRegistry()

        @registry.register
        @tool
        def get_weather(city: str) -> str:
            ...

        # Or register after decoration:
        registry.add(get_weather)

        # List all tools:
        for defn in registry.definitions():
            print(defn.name, defn.parameters)

        # Execute:
        result = await registry.execute("get_weather", {"city": "NYC"})
    """

    def __init__(self) -> None:
        self._tools: dict[str, Callable[..., Any]] = {}

    def add(self, func: Callable[..., Any]) -> Callable[..., Any]:
        """Register a @tool-decorated function."""
        defn = getattr(func, "__tool_definition__", None)
        if defn is None:
            msg = f"Function '{func.__name__}' is not decorated with @tool"
            raise ValueError(msg)
        self._tools[defn.name] = func
        return func

    def register(self, func: Callable[..., Any]) -> Callable[..., Any]:
        """Decorator form of add() — can be stacked with @tool."""
        return self.add(func)

    def get(self, name: str) -> Callable[..., Any] | None:
        """Look up a tool function by name."""
        return self._tools.get(name)

    def definitions(self) -> list[ToolDefinition]:
        """Return ToolDefinitions for all registered tools."""
        return [fn.__tool_definition__ for fn in self._tools.values()]  # type: ignore[attr-defined]

    def to_openai_tools(self) -> list[dict[str, Any]]:
        """Return tool schemas in OpenAI function-calling format."""
        return [
            {
                "type": "function",
                "function": {
                    "name": defn.name,
                    "description": defn.description,
                    "parameters": defn.parameters,
                },
            }
            for defn in self.definitions()
        ]

    async def execute(
        self,
        name: str,
        arguments: dict[str, Any] | str,
        *,
        context: Any | None = None,
    ) -> ToolResult:
        """Look up and execute a tool by name.

        Args:
            name: Tool name.
            arguments: Arguments dict or JSON string.
            context: Optional RunContext for injection.

        Returns:
            ToolResult with output or error.

        Raises:
            KeyError: If tool name is not registered.
            ToolValidationError: If arguments fail validation.
            ToolExecutionError: If the tool function raises.
        """
        func = self._tools.get(name)
        if func is None:
            raise KeyError(f"Tool not found: '{name}'")
        return await execute_tool(func, arguments, context=context)

    def __len__(self) -> int:
        return len(self._tools)

    def __contains__(self, name: str) -> bool:
        return name in self._tools
