"""@tool decorator — introspect function signatures to build JSON schemas."""

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

from pydantic import BaseModel, create_model
from pydantic.fields import FieldInfo

from koan.tools.types import ToolDefinition

# Sentinel type for RunContext injection
_CONTEXT_TYPES: set[str] = {"RunContext"}


def _is_context_param(annotation: Any) -> bool:
    """Check if a parameter annotation is a RunContext type."""
    if isinstance(annotation, str):
        return annotation in _CONTEXT_TYPES
    name = getattr(annotation, "__name__", "") or getattr(annotation, "__qualname__", "")
    return name in _CONTEXT_TYPES


def _resolve_hints(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:
        # Fall back to raw annotations (handles locally-defined types with PEP 563)
        return {k: v for k, v in func.__annotations__.items() if k != "return"}


def _resolve_annotation(annotation: Any, func: Callable[..., Any]) -> Any:
    """Resolve a string annotation to the actual type if possible."""
    if not isinstance(annotation, str):
        return annotation
    # Try to resolve from the function's globals/locals
    ns: dict[str, Any] = {}
    ns.update(getattr(func, "__globals__", {}))
    try:
        return eval(annotation, ns)  # noqa: S307
    except Exception:
        return Any


def _build_parameters_model(func: Callable[..., Any]) -> type[BaseModel]:
    """Build a Pydantic model from a function's parameters."""
    sig = inspect.signature(func)
    hints = _resolve_hints(func)

    fields: dict[str, Any] = {}
    for param_name, param in sig.parameters.items():
        if param_name == "self":
            continue

        raw_annotation = hints.get(param_name, Any)

        # Skip RunContext parameters — they're injected, not user-supplied
        if _is_context_param(raw_annotation):
            continue

        # Skip **kwargs and *args
        if param.kind in (param.VAR_POSITIONAL, param.VAR_KEYWORD):
            continue

        # Resolve string annotations to actual types for Pydantic
        annotation = _resolve_annotation(raw_annotation, func)

        # Build field info
        if isinstance(param.default, FieldInfo) or param.default is not param.empty:
            fields[param_name] = (annotation, param.default)
        else:
            fields[param_name] = (annotation, ...)

    return create_model(f"{func.__name__}_params", **fields)


def _build_json_schema(model: type[BaseModel]) -> dict[str, Any]:
    """Extract a clean JSON schema from a Pydantic model."""
    schema = model.model_json_schema()
    # Remove the title (not needed for tool schemas)
    schema.pop("title", None)
    return schema


def tool(
    func: Callable[..., Any] | None = None,
    *,
    name: str | None = None,
    description: str | None = None,
) -> Any:
    """Decorator to mark a function as a tool.

    Introspects the function signature and type hints to build a JSON schema
    for the tool's parameters. Supports both sync and async functions.

    Usage::

        @tool
        def get_weather(city: str, units: str = "celsius") -> str:
            \"\"\"Get the current weather for a city.\"\"\"
            ...

        @tool(name="search", description="Search the web")
        async def web_search(query: str, max_results: int = 10) -> list[dict]:
            ...

    The decorated function gains a ``__tool_definition__`` attribute with
    the ``ToolDefinition`` containing the JSON schema.
    """

    def decorator(fn: Callable[..., Any]) -> Callable[..., Any]:
        tool_name = name or fn.__name__
        tool_desc = description or (inspect.getdoc(fn) or "").strip() or f"Tool: {tool_name}"

        params_model = _build_parameters_model(fn)
        params_schema = _build_json_schema(params_model)

        defn = ToolDefinition(
            name=tool_name,
            description=tool_desc,
            parameters=params_schema,
            is_async=asyncio.iscoroutinefunction(fn),
        )

        fn.__tool_definition__ = defn  # type: ignore[attr-defined]
        fn.__tool_params_model__ = params_model  # type: ignore[attr-defined]
        return fn

    if func is not None:
        return decorator(func)
    return decorator
