"""Tool execution — validate inputs, call function, serialize output."""

from __future__ import annotations

import asyncio
import inspect
from typing import TYPE_CHECKING, Any, get_type_hints

if TYPE_CHECKING:
    from collections.abc import Callable

import orjson
from pydantic import BaseModel, ValidationError

from koan.tools.types import ToolExecutionError, ToolResult, ToolValidationError


def _is_context_param(annotation: Any) -> bool:
    if isinstance(annotation, str):
        return annotation in {"RunContext"}
    name = getattr(annotation, "__name__", "") or getattr(annotation, "__qualname__", "")
    return name in {"RunContext"}


def _resolve_hints_safe(func: Callable[..., Any]) -> dict[str, Any]:
    """Resolve type hints, falling back to raw annotations on failure."""
    try:
        return get_type_hints(func, include_extras=True)
    except Exception:
        return {k: v for k, v in func.__annotations__.items() if k != "return"}


async def execute_tool(
    func: Callable[..., Any],
    arguments: dict[str, Any] | str,
    *,
    context: Any | None = None,
) -> ToolResult:
    """Validate arguments, call the tool function, and return a ToolResult."""
    tool_defn = getattr(func, "__tool_definition__", None)
    tool_name: str = tool_defn.name if tool_defn else func.__name__
    params_model: type[BaseModel] | None = getattr(func, "__tool_params_model__", None)

    # Parse JSON string arguments into a dict
    if isinstance(arguments, str):
        try:
            args: dict[str, Any] = orjson.loads(arguments)
        except (orjson.JSONDecodeError, ValueError) as e:
            return ToolResult(tool_name=tool_name, error=f"Invalid JSON arguments: {e}")
    else:
        args = arguments

    # Validate with Pydantic model if available
    if params_model is not None:
        try:
            validated = params_model.model_validate(args)
            args = validated.model_dump()
        except ValidationError as e:
            raise ToolValidationError(
                f"Validation failed for tool '{tool_name}': {e}",
                tool_name=tool_name,
                cause=e,
            ) from e

    # Inject RunContext if the function expects it
    if context is not None:
        hints = _resolve_hints_safe(func)
        sig = inspect.signature(func)
        for param_name in sig.parameters:
            annotation = hints.get(param_name)
            if annotation is not None and _is_context_param(annotation):
                args[param_name] = context

    # Execute the function
    try:
        if asyncio.iscoroutinefunction(func):
            result = await func(**args)
        else:
            result = func(**args)
    except Exception as e:
        raise ToolExecutionError(
            f"Tool '{tool_name}' raised: {e}",
            tool_name=tool_name,
            cause=e,
        ) from e

    output = result.model_dump() if isinstance(result, BaseModel) else result
    return ToolResult(tool_name=tool_name, output=output)
