"""JSON Schema generation utilities."""

from __future__ import annotations

import inspect
from collections.abc import Callable
from typing import Any, get_args, get_origin, get_type_hints

from pydantic import BaseModel


def function_schema(fn: Callable[..., Any]) -> dict[str, Any]:
    """Generate a JSON Schema for a function's input from its signature.

    Supports:
    - Plain Python type hints (str, int, float, bool, list, dict, etc.)
    - Pydantic BaseModel subclasses as root input
    - Optional[X] / Union[X, None]
    - Default values from the signature

    Returns a JSON Schema dict with type=object, properties, and required.
    """
    sig = inspect.signature(fn)
    hints = get_type_hints(fn)

    # If the return / first param is a Pydantic model, use its schema as root.
    for param_name, param in sig.parameters.items():
        ann = hints.get(param_name, param.annotation)
        if ann is inspect.Parameter.empty:
            continue
        if isinstance(ann, type) and issubclass(ann, BaseModel):
            schema = ann.model_json_schema()
            schema.pop("title", None)
            return _strip_unsupported(schema)

    properties: dict[str, Any] = {}
    required: list[str] = []

    for param_name, param in sig.parameters.items():
        if param_name == "self":
            continue
        ann = hints.get(param_name, param.annotation)
        if ann is inspect.Parameter.empty:
            continue
        properties[param_name] = _annotation_to_schema(ann)
        if param.default is inspect.Parameter.empty:
            required.append(param_name)

    return {"type": "object", "properties": properties, "required": required, "additionalProperties": False}


def return_schema(fn: Callable[..., Any]) -> dict[str, Any]:
    """Generate JSON Schema for a function's return type."""
    hints = get_type_hints(fn)
    ret = hints.get("return", inspect.signature(fn).return_annotation)
    if ret is inspect.Signature.empty:
        return {"type": "object", "additionalProperties": True}
    if isinstance(ret, type) and issubclass(ret, BaseModel):
        schema = ret.model_json_schema()
        schema.pop("title", None)
        return _strip_unsupported(schema)
    return _annotation_to_schema(ret)


def _annotation_to_schema(ann: Any) -> dict[str, Any]:
    origin = get_origin(ann)
    args = get_args(ann)

    # Plain types
    if ann is str:
        return {"type": "string"}
    if ann is int:
        return {"type": "integer"}
    if ann is float:
        return {"type": "number"}
    if ann is bool:
        return {"type": "boolean"}

    # Containers — check BEFORE generic Union handling, since
    # dict[str, X] / list[X] carry args too.
    if origin is dict or ann is dict:
        if args and len(args) == 2:
            return {
                "type": "object",
                "additionalProperties": _annotation_to_schema(args[1]),
            }
        return {"type": "object", "additionalProperties": True}
    if origin is list or ann is list:
        if args and len(args) == 1:
            return {"type": "array", "items": _annotation_to_schema(args[0])}
        return {"type": "array"}

    # Optional[X] / Union[X, None] / Union[X, Y]
    if origin is not None and args:
        non_none = [a for a in args if a is not type(None)]
        if len(non_none) == 1 and len(args) > len(non_none):
            sub = _annotation_to_schema(non_none[0])
            sub["nullable"] = True
            return sub
        if len(non_none) > 1:
            return {"anyOf": [_annotation_to_schema(a) for a in non_none]}

    if isinstance(ann, type) and issubclass(ann, BaseModel):
        schema = ann.model_json_schema()
        schema.pop("title", None)
        return _strip_unsupported(schema)

    return {"type": "object", "additionalProperties": True}


def _strip_unsupported(schema: dict[str, Any]) -> dict[str, Any]:
    """Remove JSON Schema keys that some LLM tool-use pipelines reject."""
    schema.pop("title", None)
    return schema