"""The `ToolSchema` class for defining tools that LLMs can request be called."""

from __future__ import annotations

import inspect
from dataclasses import dataclass
from typing import (
    Annotated,
    Any,
    Generic,
    NamedTuple,
    TypeAlias,
    TypedDict,
    TypeVar,
    cast,
    get_args,
    get_origin,
    get_type_hints,
)

from docstring_parser import parse
from pydantic import BaseModel, Field, create_model
from pydantic.fields import FieldInfo

from ..._utils import copy_function_metadata
from ..content import ToolCall
from ..types import Jsonable
from .protocols import AsyncContextToolFn, AsyncToolFn, ContextToolFn, ToolFn

AnyToolFn: TypeAlias = (
    ToolFn[..., Jsonable]
    | AsyncToolFn[..., Jsonable]
    | ContextToolFn[Any, ..., Jsonable]
    | AsyncContextToolFn[Any, ..., Jsonable]
)

ToolFnT = TypeVar(
    "ToolFnT",
    bound=AnyToolFn,
    covariant=True,
)

ToolSchemaT = TypeVar("ToolSchemaT", bound="ToolSchema[AnyToolFn]", covariant=True)


ModelJsonSchema = TypedDict(
    "ModelJsonSchema",
    {
        "properties": dict[str, dict[str, Any]],
        "required": list[str],
        "$defs": dict[str, dict[str, Any]],
    },
    total=False,
)
"""Type for Pydantic's model_json_schema() output.

This TypedDict defines the structure of JSON schemas returned by Pydantic models,
allowing us to avoid type casts when extracting schema components.
"""


class DocstringArg(NamedTuple):
    """A parameter from a docstring."""

    name: str
    description: str


FORMAT_TOOL_NAME = "__mirascope_formatted_output_tool__"
"""Reserved name of the formatted output tool. 

Any call to a tool with this name is NOT considered a regular tool call, but will instead
be converted into textual output containing the arguments to the tool call.
"""
# TODO: Investigate whether using this hardcoded tool name has any adverse impact
# on model performance, compared to a tool name that references the name of the formatted
# class.


@dataclass
class ParsedDocstring:
    args: list[DocstringArg]


def _parse_docstring_params(docstring: str | None) -> ParsedDocstring:
    """Parse parameter descriptions from a docstring.

    Uses docstring-parser library which supports ReST, Google, Numpydoc-style and
    Epydoc docstrings

    Args:
        docstring: The function's docstring

    Returns:
        ParsedDocstring containing parameter descriptions
    """
    if not docstring:
        return ParsedDocstring(args=[])

    parsed = parse(docstring)
    args: list[DocstringArg] = []

    for param in parsed.params:
        if param.description:
            args.append(
                DocstringArg(name=param.arg_name, description=param.description)
            )

    return ParsedDocstring(args=args)


class ToolParameterSchema(BaseModel):
    """JSON Schema for tool parameters (always an object with properties).

    This contains real JSON Schema as generated by Pydantic, with full support
    for complex schemas like anyOf, nested objects, validation constraints, etc.
    Including $defs for complex type references.
    """

    properties: dict[str, dict[str, Any]] = Field(default_factory=dict)
    """Dictionary mapping parameter names to their JSON Schema definitions."""

    required: list[str] = Field(default_factory=list)
    """List of required parameter names."""

    additionalProperties: bool = False
    """Whether additional properties beyond those defined are allowed."""

    defs: dict[str, dict[str, Any]] | None = Field(default=None, alias="$defs")
    """JSON Schema definitions for complex types referenced via $ref."""


@dataclass
class ToolSchema(Generic[ToolFnT]):
    """Underlying schema defining a tool that can be used by LLMs.

    A ToolSchema represents a function that can be called by an LLM during a call.
    It includes metadata like name, description, and parameter schema.

    This class is not instantiated directly but created by the `@tool()` decorator.
    """

    fn: ToolFnT
    """The function that implements the tool's functionality."""

    name: str
    """The name of the tool, used by the LLM to identify which tool to call."""

    description: str
    """Description of what the tool does, extracted from the function's docstring."""

    parameters: ToolParameterSchema
    """JSON Schema describing the parameters accepted by the tool.
    
    The serialized parameters table is cached for efficient hash lookups (e.g. when
    caching provider-encoded tool representations in a LRU cache). Therefore,
    it should **not be modified** after the ToolSchema is created.
    """

    strict: bool | None
    """Whether the tool should use strict mode when supported by the model.

    If set to None, will use the provider's default setting (usually as strict as
    possible).
    """

    __name__: str
    """The name of the underlying function (preserved for decorator stacking)."""

    def __hash__(self) -> int:
        if not hasattr(self, "_hash"):
            self._hash = hash(
                (
                    self.name,
                    self.description,
                    self.strict,
                    self.parameters.model_dump_json(),
                )
            )
        return self._hash

    def __init__(
        self,
        fn: ToolFnT,
        name: str,
        description: str,
        parameters: ToolParameterSchema,
        *,
        strict: bool | None = None,
    ) -> None:
        """Create a `ToolSchema` with the provided values.

        Args:
            fn: The function that implements the tool's functionality
            name: The name of the tool
            description: Description of what the tool does
            parameters: JSON Schema describing the parameters accepted by the tool
            strict: Whether the tool should use strict mode when supported.
                If None, uses provider's default (usually as strict as possible).
        """
        self.fn = fn
        self.name = name
        self.description = description
        self.parameters = parameters
        self.strict = strict
        copy_function_metadata(self, fn)

    @classmethod
    def from_function(
        cls,
        fn: AnyToolFn,
        *,
        strict: bool | None = None,
        is_context_tool: bool = False,
    ) -> ToolSchema[AnyToolFn]:
        """Create a `ToolSchema` by inspecting a function and its docstring.

        Uses Pydantic's create_model to dynamically build a model from the function
        signature, then extracts the JSON schema from it. This leverages Pydantic's
        robust type conversion and union handling.

        Args:
            fn: The function to extract schema from
            strict: Whether the tool should use strict mode when supported.
                If None, uses provider's default (usually as strict as possible).
            is_context_tool: Whether this is a context tool (skips the context parameter)

        Returns:
            a `ToolSchema` representing the function

        Raises:
            ValueError: If the tool has a reserved name.
        """
        name = fn.__name__
        if name.startswith(FORMAT_TOOL_NAME):
            raise ValueError(
                f"Cannot use reserved name {FORMAT_TOOL_NAME} as tool name."
            )
        description = inspect.cleandoc(fn.__doc__) if fn.__doc__ else name

        param_descriptions = _parse_docstring_params(fn.__doc__)

        field_definitions: dict[str, tuple[Any, Any]] = {}
        hints = get_type_hints(fn, include_extras=True)

        context_param_skipped = False
        for param in inspect.signature(fn).parameters.values():
            # Skip args that are provided by Python or Mirascope, not LLM-generated.
            # TODO: Handling of FromCallArgs
            if param.name in ("self", "cls"):
                continue

            # For context tools, skip the first non-self/cls parameter (the context parameter)
            if is_context_tool and not context_param_skipped:
                context_param_skipped = True
                continue

            param_type = hints.get(param.name, Any)
            default = ... if param.default is inspect.Parameter.empty else param.default
            field_info = None

            if get_origin(param_type) is Annotated:
                args = get_args(param_type)
                param_type = args[0]
                for annotation in args[1:]:
                    if isinstance(annotation, FieldInfo):
                        field_info = annotation
                        break

            if field_info is not None:
                field_value = Field(
                    default=default,
                    description=field_info.description,
                )
            else:
                docstring_description: str | None = None
                for arg in param_descriptions.args:
                    if arg.name == param.name:
                        docstring_description = arg.description
                        break
                if docstring_description:
                    field_value = Field(
                        default=default, description=docstring_description
                    )
                else:
                    field_value = default

            field_definitions[param.name] = (param_type, field_value)

        TempModel = create_model("TempModel", **cast(dict[str, Any], field_definitions))

        schema = cast(ModelJsonSchema, TempModel.model_json_schema())

        parameters = ToolParameterSchema(
            properties=schema.get("properties", {}),
            required=schema.get("required", []),
            additionalProperties=False,
        )
        if "$defs" in schema:
            parameters.defs = schema["$defs"]

        return cls(
            fn=cast(ToolFnT, fn),
            name=name,
            description=description,
            parameters=parameters,
            strict=strict,
        )

    def can_execute(self, tool_call: ToolCall) -> bool:
        """Check if a `ToolCall` can be executed by tools with this `ToolSchema`.

        This method is a convenient way to determine if a `ToolCall` is likely intended
        to be executed by a tool with this `ToolSchema`. It does so by checking
        whether the name on the call matches the name on the schema. No other validation
        is performed.
        """
        return tool_call.name == self.name


AnyToolSchema: TypeAlias = ToolSchema[AnyToolFn]
