from collections.abc import Callable
from inspect import _empty, signature
from typing import Any, get_type_hints

from pydantic import (
    Field,
    PydanticSchemaGenerationError,
    create_model,
)

from ixoncdkingress.types import Schema, ToolIssue

# ---------------------------------------------------------------------------
# Parameter collection
# ---------------------------------------------------------------------------


def _build_input_fields(
    sig: Any,
    type_hints: dict[str, Any],
    *,
    exclude_param_types: set[type] | None = None,
    exclude_param_names: set[str] | None = None,
) -> dict[str, Any]:
    """Return ``{name: (hint, default)}`` for non-excluded parameters."""
    if exclude_param_types is None:
        exclude_param_types = set()

    if exclude_param_names is None:
        exclude_param_names = set()

    fields: dict[str, Any] = {}
    for name, param in sig.parameters.items():
        if param.annotation is _empty:
            continue
        hint = type_hints.get(name)
        if hint in exclude_param_types or name in exclude_param_names or name.startswith("_"):
            continue
        default = param.default if param.default is not _empty else ...
        fields[name] = (hint, default)
    return fields


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------


def generate_input_schema(
    function: Callable[..., Any],
    *,
    exclude_param_types: set[type] | None = None,
    exclude_param_names: set[str] | None = None,
) -> Schema:
    """Generate the input schema for a function.

    Uses pydantic's `create_model` / `model_json_schema` for the base schema.

    Annotated aliases with `Field(json_schema_extra={"definition": ...})` are propagated natively
    by pydantic for all positions(union members, list/set/tuple items, dict keys) via
    `propertyNames` and for dict values this via `additionalProperties`.
    """
    sig = signature(function)
    # include_extras=True preserves Annotated metadata so pydantic can merge
    # json_schema_extra into the generated schema nodes.
    type_hints = get_type_hints(function, include_extras=True)
    merged_exclude_names: set[str] = set(exclude_param_names) if exclude_param_names else set()

    fields = _build_input_fields(
        sig,
        type_hints,
        exclude_param_types=exclude_param_types,
        exclude_param_names=merged_exclude_names,
    )
    try:
        input_model = create_model("InputSchema", **fields)
        input_schema = input_model.model_json_schema()
        input_schema.pop("title", None)
        return Schema(
            schema=input_schema,
            model=input_model,
        )
    except PydanticSchemaGenerationError as e:
        return Schema(
            issue=ToolIssue(
                type="input_schema_generation_error",
                error=str(e),
            ),
        )


def generate_output_schema(function: Callable[..., Any]) -> Schema:
    """
    Generate the output schema for a function.

    Uses pydantic's `create_model` / `model_json_schema` for the base schema.

    Annotated aliases with `Field(json_schema_extra={"definition": ...})` are propagated natively
    by pydantic for all positions(union members, list/set/tuple items, dict keys) via
    `propertyNames` and for dict values this via `additionalProperties`.
    """
    sig = signature(function)
    # include_extras=True preserves Annotated metadata so pydantic can merge
    # json_schema_extra into the generated schema nodes.
    type_hints = get_type_hints(function, include_extras=True)

    if sig.return_annotation is _empty:
        return Schema()
    try:
        output_model = create_model(
            "OutputSchema",
            output=(type_hints["return"], Field(default=...)),
        )
        output_schema = output_model.model_json_schema()
        output_schema.pop("title")
        return Schema(
            schema=output_schema,
            model=output_model,
        )

    except PydanticSchemaGenerationError as e:
        return Schema(
            issue=ToolIssue(
                type="output_schema_generation_error",
                error=str(e),
            ),
        )
