"""Extract JSON Schema from Python type hints and Pydantic models.

Supports: str, int, float, bool, list[T], dict, Optional[T], T | None,
Literal["a","b"], Annotated[T, ...], and Pydantic BaseModel.
"""
from __future__ import annotations

import inspect
import sys
from typing import Any, Union, get_args, get_origin, get_type_hints

from pillar._context import ToolContext

_PRIMITIVE_MAP: dict[type, str] = {
    str: "string",
    int: "integer",
    float: "number",
    bool: "boolean",
}

_EXCLUDED_PARAMS = frozenset({"self", "cls", "ctx", "context", "return"})


def _is_optional(annotation: Any) -> tuple[bool, Any]:
    """Return (True, inner_type) if *annotation* is Optional[T] or T | None."""
    import types

    origin = get_origin(annotation)
    # typing.Optional[T] / typing.Union[T, None]
    if origin is Union:
        args = get_args(annotation)
        non_none = [a for a in args if a is not type(None)]
        if len(non_none) == 1 and type(None) in args:
            return True, non_none[0]
    # PEP 604: T | None  (Python 3.10+ produces types.UnionType)
    if isinstance(annotation, types.UnionType):
        args = get_args(annotation)
        non_none = [a for a in args if a is not type(None)]
        if len(non_none) == 1 and type(None) in args:
            return True, non_none[0]
    return False, annotation


def _is_literal(annotation: Any) -> tuple[bool, tuple[Any, ...]]:
    """Return (True, values) if *annotation* is Literal[...]."""
    import typing
    origin = get_origin(annotation)
    if origin is getattr(typing, "Literal", None):
        return True, get_args(annotation)
    return False, ()


def _is_annotated(annotation: Any) -> tuple[bool, Any, tuple[Any, ...]]:
    """Return (True, inner_type, metadata) if *annotation* is Annotated[T, ...]."""
    if get_origin(annotation) is getattr(__import__("typing"), "Annotated", None):
        args = get_args(annotation)
        return True, args[0], args[1:]
    # Python 3.11+ has typing.Annotated; 3.9/3.10 may use typing_extensions
    if sys.version_info < (3, 11):
        try:
            from typing_extensions import Annotated as AnnotatedExt, get_origin as ge_ext
            if ge_ext(annotation) is AnnotatedExt:
                args = get_args(annotation)
                return True, args[0], args[1:]
        except ImportError:
            pass
    return False, annotation, ()


def _extract_description_from_metadata(metadata: tuple[Any, ...]) -> str | None:
    """Pull a description string from Annotated metadata.

    Supports plain strings, objects with a ``.description`` attribute (like
    Pydantic ``Field``), and dicts with a ``"description"`` key.
    """
    for item in metadata:
        if isinstance(item, str):
            return item
        if hasattr(item, "description") and isinstance(item.description, str):
            return item.description
        if isinstance(item, dict) and "description" in item:
            return item["description"]
    return None


def _type_to_schema(annotation: Any) -> dict[str, Any]:
    """Convert a single Python type annotation to a JSON Schema fragment."""

    # Annotated[T, ...]
    is_ann, inner, metadata = _is_annotated(annotation)
    if is_ann:
        schema = _type_to_schema(inner)
        desc = _extract_description_from_metadata(metadata)
        if desc:
            schema["description"] = desc
        return schema

    # Optional[T] / T | None  — unwrap, schema is for the inner type
    is_opt, inner = _is_optional(annotation)
    if is_opt:
        return _type_to_schema(inner)

    # Literal["a", "b", ...]
    is_lit, values = _is_literal(annotation)
    if is_lit:
        types = {type(v) for v in values}
        json_type = _PRIMITIVE_MAP.get(types.pop(), "string") if len(types) == 1 else "string"
        return {"type": json_type, "enum": list(values)}

    # Primitives
    if annotation in _PRIMITIVE_MAP:
        return {"type": _PRIMITIVE_MAP[annotation]}

    # list / list[T]
    origin = get_origin(annotation)
    if origin is list:
        args = get_args(annotation)
        schema: dict[str, Any] = {"type": "array"}
        if args:
            schema["items"] = _type_to_schema(args[0])
        return schema

    # dict / dict[K, V]
    if origin is dict or annotation is dict:
        return {"type": "object"}

    # Pydantic BaseModel
    try:
        from pydantic import BaseModel
        if isinstance(annotation, type) and issubclass(annotation, BaseModel):
            return annotation.model_json_schema()
    except ImportError:
        pass

    # Fallback
    return {"type": "string"}


def _should_exclude_param(name: str, annotation: Any) -> bool:
    """Return True if this parameter should be excluded from the schema."""
    if name in _EXCLUDED_PARAMS:
        return True
    if annotation is ToolContext:
        return True
    if isinstance(annotation, type) and issubclass(annotation, ToolContext):
        return True
    return False


def extract_schema(func: Any) -> dict[str, Any]:
    """Extract a JSON Schema ``object`` from a function's signature.

    Returns a dict like::

        {
            "type": "object",
            "properties": { ... },
            "required": ["param1", "param2"]
        }
    """
    sig = inspect.signature(func)
    try:
        hints = get_type_hints(func, include_extras=True)
    except Exception:
        hints = {}

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

    for name, param in sig.parameters.items():
        annotation = hints.get(name, param.annotation)
        if annotation is inspect.Parameter.empty:
            annotation = str

        if _should_exclude_param(name, annotation):
            continue

        prop_schema = _type_to_schema(annotation)
        properties[name] = prop_schema

        is_opt, _ = _is_optional(annotation)
        has_default = param.default is not inspect.Parameter.empty
        if not is_opt and not has_default:
            required.append(name)

    schema: dict[str, Any] = {"type": "object", "properties": properties}
    if required:
        schema["required"] = required
    return schema
