"""Function introspection and schema generation for FastMCP tools."""

from __future__ import annotations

import functools
import inspect
import types
from collections.abc import Callable
from dataclasses import dataclass
from typing import Annotated, Any, Generic, Union, get_args, get_origin, get_type_hints

import mcp.types
from pydantic import BaseModel, PydanticSchemaGenerationError
from typing_extensions import TypeVar as TypeVarExt

from fastmcp.tools.base import ToolResult, resolve_serialize_by_alias
from fastmcp.utilities.docstring_parsing import ParsedDocstring, parse_docstring
from fastmcp.utilities.json_schema import compress_schema
from fastmcp.utilities.logging import get_logger
from fastmcp.utilities.types import (
    Audio,
    File,
    Image,
    create_function_without_params,
    get_cached_typeadapter,
    is_class_member_of_type,
    replace_type,
)

try:
    from prefab_ui.app import PrefabApp as _PrefabApp
    from prefab_ui.components.base import Component as _PrefabComponent

    _PREFAB_TYPES: tuple[type, ...] = (_PrefabApp, _PrefabComponent)
except ImportError:
    _PREFAB_TYPES = ()


def _contains_bytes_type(tp: Any) -> bool:
    """Check if *tp* is or contains bytes, recursing through unions and Annotated."""
    if tp is bytes:
        return True
    origin = get_origin(tp)
    if origin is Union or origin is types.UnionType or origin is Annotated:
        return any(_contains_bytes_type(a) for a in get_args(tp))
    return False


def _contains_prefab_type(tp: Any) -> bool:
    """Check if *tp* is or contains a prefab type, recursing through unions and Annotated."""
    if isinstance(tp, type) and issubclass(tp, _PREFAB_TYPES):
        return True
    origin = get_origin(tp)
    if origin is Union or origin is types.UnionType or origin is Annotated:
        return any(_contains_prefab_type(a) for a in get_args(tp))
    return False


def _unwrap_model(tp: Any) -> type[BaseModel] | None:
    """Unwrap ``Annotated`` and return the underlying Pydantic model, if any."""
    if get_origin(tp) is Annotated:
        return _unwrap_model(get_args(tp)[0])
    if isinstance(tp, type) and issubclass(tp, BaseModel):
        return tp
    return None


def _resolve_output_by_alias(tp: Any) -> bool:
    """Resolve ``by_alias`` for the output schema of return type *tp*.

    Unwraps ``Annotated`` and ``Optional``/``Union`` wrappers to find the
    underlying Pydantic model so the generated schema honors the model's
    ``serialize_by_alias`` config — keeping it consistent with how the runtime
    result is serialized. Containers (``list[Model]`` etc.) are not unwrapped:
    their schema keeps the default, matching the runtime path which only
    special-cases a directly-returned model.

    Known limitation: a single schema is generated with one ``by_alias`` value,
    while the runtime resolves the alias mode per returned value. They cannot
    diverge for a plain single-model return, but a union return can produce more
    than one runtime alias mode that no single schema can describe:

    - distinct models with *conflicting* ``serialize_by_alias`` (e.g. ``A | B``
      where ``A`` opts out but ``B`` opts in), and
    - a model arm alongside a container arm (e.g. ``Model | list[Model]``):
      a directly-returned model honors its config, but a returned ``list`` is
      serialized with the default alias mode, so the two variants disagree.

    Pydantic's schema generator does not consult per-model ``serialize_by_alias``
    and the runtime does not recurse into containers, so honoring every variant
    would require per-arm schema assembly. This is an accepted edge; single-model
    returns and unions whose arms all resolve to the same mode are consistent.
    """
    origin = get_origin(tp)
    if origin is Annotated:
        return _resolve_output_by_alias(get_args(tp)[0])
    if origin is Union or origin is types.UnionType:
        for arg in get_args(tp):
            model = _unwrap_model(arg)
            if model is not None:
                return resolve_serialize_by_alias(model)
        return True
    return resolve_serialize_by_alias(tp)


T = TypeVarExt("T", default=Any)

logger = get_logger(__name__)


@dataclass
class _WrappedResult(Generic[T]):
    result: T


class _UnserializableType:
    pass


def _is_object_schema(
    schema: dict[str, Any],
    *,
    _root_schema: dict[str, Any] | None = None,
    _seen_refs: set[str] | None = None,
) -> bool:
    """Check if a JSON schema represents an object type."""
    root_schema = _root_schema or schema
    seen_refs = _seen_refs or set()

    # Direct object type
    if schema.get("type") == "object":
        return True

    # Schema with properties but no explicit type is treated as object
    if "properties" in schema:
        return True

    # Resolve local $ref definitions and recurse into the target schema.
    ref = schema.get("$ref")
    if not isinstance(ref, str) or not ref.startswith("#/"):
        return False

    if ref in seen_refs:
        return False

    # Walk the JSON Pointer path from the root schema, unescaping each
    # token per RFC 6901 (~1 → /, ~0 → ~).
    pointer = ref.removeprefix("#/")
    segments = pointer.split("/")
    target: Any = root_schema
    for segment in segments:
        unescaped = segment.replace("~1", "/").replace("~0", "~")
        if not isinstance(target, dict) or unescaped not in target:
            return False
        target = target[unescaped]

    target_schema = target
    if not isinstance(target_schema, dict):
        return False

    return _is_object_schema(
        target_schema,
        _root_schema=root_schema,
        _seen_refs=seen_refs | {ref},
    )


@dataclass
class ParsedFunction:
    fn: Callable[..., Any]
    name: str
    description: str | None
    input_schema: dict[str, Any]
    output_schema: dict[str, Any] | None
    return_type: Any = None

    @classmethod
    def from_function(
        cls,
        fn: Callable[..., Any],
        exclude_args: list[str] | None = None,
        validate: bool = True,
        wrap_non_object_output_schema: bool = True,
    ) -> ParsedFunction:
        if validate:
            sig = inspect.signature(fn)
            # Reject functions with *args or **kwargs
            for param in sig.parameters.values():
                if param.kind == inspect.Parameter.VAR_POSITIONAL:
                    raise ValueError("Functions with *args are not supported as tools")
                if param.kind == inspect.Parameter.VAR_KEYWORD:
                    raise ValueError(
                        "Functions with **kwargs are not supported as tools"
                    )

            # Reject exclude_args that don't exist in the function or don't have a default value
            if exclude_args:
                for arg_name in exclude_args:
                    if arg_name not in sig.parameters:
                        raise ValueError(
                            f"Parameter '{arg_name}' in exclude_args does not exist in function."
                        )
                    param = sig.parameters[arg_name]
                    if param.default == inspect.Parameter.empty:
                        raise ValueError(
                            f"Parameter '{arg_name}' in exclude_args must have a default value."
                        )

        # collect name and description before we potentially modify the function
        fn_name = getattr(fn, "__name__", None) or fn.__class__.__name__
        outer_docstring = parse_docstring(fn)

        # if the fn is a callable class, we need to get the __call__ method from here out
        if not inspect.isroutine(fn) and not isinstance(fn, functools.partial):
            fn = fn.__call__
        # if the fn is a staticmethod, we need to work with the underlying function
        if isinstance(fn, staticmethod):
            fn = fn.__func__

        # For callable classes, parameter descriptions must come from
        # __call__'s docstring — where the exposed parameters are actually
        # declared. The class docstring's Args section, if any, typically
        # describes __init__, so falling back to it would risk injecting
        # constructor docs into __call__'s schema on overlapping names.
        # The description, however, comes from the class docstring (which
        # describes what the tool IS) when present.
        inner_docstring = parse_docstring(fn)
        parsed_docstring = ParsedDocstring(
            description=outer_docstring.description or inner_docstring.description,
            parameters=inner_docstring.parameters,
        )

        # Transform Context type annotations to Depends() for unified DI
        from fastmcp.server.dependencies import (
            transform_context_annotations,
            without_injected_parameters,
        )

        fn = transform_context_annotations(fn)

        # Handle injected parameters (Context, Docket dependencies)
        wrapper_fn = without_injected_parameters(fn)

        # Also handle exclude_args with non-serializable types (issue #2431)
        # This must happen before Pydantic tries to serialize the parameters
        if exclude_args:
            wrapper_fn = create_function_without_params(wrapper_fn, list(exclude_args))

        input_type_adapter = get_cached_typeadapter(wrapper_fn)
        input_schema = input_type_adapter.json_schema()

        # Compress and handle exclude_args
        prune_params = list(exclude_args) if exclude_args else None
        input_schema = compress_schema(
            input_schema, prune_params=prune_params, prune_titles=True
        )

        # Inject parameter descriptions from the docstring into the schema.
        # Explicit annotations (Field(description=...), Annotated[x, "..."])
        # already have a "description" key and take precedence.
        if parsed_docstring.parameters:
            properties = input_schema.get("properties", {})
            for param_name, param_desc in parsed_docstring.parameters.items():
                if (
                    param_name in properties
                    and "description" not in properties[param_name]
                ):
                    properties[param_name]["description"] = param_desc

        output_schema = None
        # Get the return annotation from the signature
        sig = inspect.signature(fn)
        output_type = sig.return_annotation

        # If the annotation is a string (from __future__ annotations), resolve it
        if isinstance(output_type, str):
            try:
                # Use get_type_hints to resolve the return type
                # include_extras=True preserves Annotated metadata
                type_hints = get_type_hints(fn, include_extras=True)
                output_type = type_hints.get("return", output_type)
            except Exception as e:
                # If resolution fails, keep the string annotation
                logger.debug("Failed to resolve type hint for return annotation: %s", e)

        # Save original for return_type before any schema-related replacement
        original_output_type = output_type

        if output_type not in (inspect._empty, None, Any, ...):
            # bytes can't be represented as structured JSON output — skip schema
            if _contains_bytes_type(output_type):
                output_type = _UnserializableType

            # Prefab component subclasses (Column, Card, etc.) shouldn't
            # produce output schemas — replace_type only does exact matching,
            # so we handle subclass matching explicitly here.  We also need
            # to handle composite types like ``Column | None`` and
            # ``Annotated[PrefabApp, ...]`` by recursing into their args.
            if _PREFAB_TYPES and _contains_prefab_type(output_type):
                output_type = _UnserializableType

            # ToolResult subclasses should suppress schema generation just
            # like ToolResult itself — replace_type only does exact matching.
            if is_class_member_of_type(output_type, ToolResult):
                output_type = _UnserializableType

            # there are a variety of types that we don't want to attempt to
            # serialize because they are either used by FastMCP internally,
            # or are MCP content types that explicitly don't form structured
            # content. By replacing them with an explicitly unserializable type,
            # we ensure that no output schema is automatically generated.
            clean_output_type = replace_type(
                output_type,
                dict.fromkeys(
                    (
                        Image,
                        Audio,
                        File,
                        ToolResult,
                        mcp.types.TextContent,
                        mcp.types.ImageContent,
                        mcp.types.AudioContent,
                        mcp.types.ResourceLink,
                        mcp.types.EmbeddedResource,
                        *_PREFAB_TYPES,
                    ),
                    _UnserializableType,
                ),
            )

            try:
                # Honor the model's serialize_by_alias config so the schema's
                # field names match the serialized result (see base.py).
                by_alias = _resolve_output_by_alias(clean_output_type)
                type_adapter = get_cached_typeadapter(clean_output_type)
                base_schema = type_adapter.json_schema(
                    mode="serialization", by_alias=by_alias
                )

                # Generate schema for wrapped type if it's non-object
                # because MCP requires that output schemas are objects
                # Check if schema is an object type, resolving $ref references
                # (self-referencing types use $ref at root level)
                if wrap_non_object_output_schema and not _is_object_schema(base_schema):
                    # Use the wrapped result schema directly
                    wrapped_type = _WrappedResult[clean_output_type]
                    wrapped_adapter = get_cached_typeadapter(wrapped_type)
                    output_schema = wrapped_adapter.json_schema(
                        mode="serialization", by_alias=by_alias
                    )
                    output_schema["x-fastmcp-wrap-result"] = True
                else:
                    output_schema = base_schema

                output_schema = compress_schema(output_schema, prune_titles=True)

            except PydanticSchemaGenerationError as e:
                if "_UnserializableType" not in str(e):
                    logger.debug(f"Unable to generate schema for type {output_type!r}")

        return cls(
            fn=fn,
            name=fn_name,
            description=parsed_docstring.description,
            input_schema=input_schema,
            output_schema=output_schema or None,
            return_type=original_output_type,
        )
