from __future__ import annotations

import inspect
from dataclasses import dataclass
from typing import Any, Dict, Optional, Tuple, Type, get_type_hints

from pydantic import BaseModel, create_model


@dataclass(frozen=True)
class ArgsModelBundle:
    """Pydantic model + OpenAI JSON-schema parameters extracted from it."""

    model: Type[BaseModel]
    parameters_schema: Dict[str, Any]


def build_args_model(func, *, model_name: Optional[str] = None) -> Type[BaseModel]:
    """Create a Pydantic model representing the callable signature."""
    sig = inspect.signature(func)
    hints = get_type_hints(func)

    fields: Dict[str, Tuple[Any, Any]] = {}

    for param in sig.parameters.values():
        if param.kind in (
            inspect.Parameter.VAR_KEYWORD,
            inspect.Parameter.VAR_POSITIONAL,
        ):
            raise ValueError(
                "Variadic parameters (*args/**kwargs) are not supported for tools"
            )

        if param.name == "self":
            raise ValueError(
                "Instance methods are not supported; pass a function or @staticmethod"
            )

        annotation = hints.get(param.name, Any)

        if param.default is inspect._empty:
            default = ...
        else:
            default = param.default

        fields[param.name] = (annotation, default)

    name = model_name or f"{func.__name__}_Args"
    return create_model(name, **fields)  # type: ignore[arg-type]


def _ensure_null_allowed(prop_schema: Dict[str, Any]) -> Dict[str, Any]:
    """Wrap a schema so that null is allowed (best-effort)."""
    if "type" in prop_schema:
        typ = prop_schema["type"]
        if isinstance(typ, list):
            if "null" not in typ:
                prop_schema = dict(prop_schema)
                prop_schema["type"] = list(typ) + ["null"]
            return prop_schema
        if isinstance(typ, str) and typ != "null":
            prop_schema = dict(prop_schema)
            prop_schema["type"] = [typ, "null"]
            return prop_schema

    if "anyOf" in prop_schema and isinstance(prop_schema["anyOf"], list):
        any_of = list(prop_schema["anyOf"])
        if not any(isinstance(x, dict) and x.get("type") == "null" for x in any_of):
            any_of.append({"type": "null"})
        prop_schema = dict(prop_schema)
        prop_schema["anyOf"] = any_of
        return prop_schema

    # Fallback: wrap as anyOf
    return {"anyOf": [prop_schema, {"type": "null"}]}


def build_parameters_schema(
    args_model: Type[BaseModel],
    *,
    strict: bool,
    force_all_required: bool = True,
    allow_null_for_optional: bool = True,
) -> Dict[str, Any]:
    """Extract OpenAI-style JSON schema for function tool parameters."""
    raw = args_model.model_json_schema()

    schema: Dict[str, Any] = {
        "type": "object",
        "properties": raw.get("properties", {}) or {},
    }

    required = list(raw.get("required", []) or [])

    if strict:
        schema["additionalProperties"] = False

        if force_all_required:
            all_keys = list(schema["properties"].keys())
            schema["required"] = all_keys

            if allow_null_for_optional:
                for key in all_keys:
                    if key in required:
                        continue
                    prop = schema["properties"].get(key)
                    if isinstance(prop, dict):
                        schema["properties"][key] = _ensure_null_allowed(prop)
        else:
            schema["required"] = required
    else:
        if required:
            schema["required"] = required

    return schema


def build_tool_schema_from_callable(
    func,
    *,
    strict: bool,
    force_all_required: bool = True,
    allow_null_for_optional: bool = True,
) -> ArgsModelBundle:
    """Build (ArgsModel, parameters_schema) from a Python callable."""
    model = build_args_model(func)
    parameters_schema = build_parameters_schema(
        model,
        strict=strict,
        force_all_required=force_all_required,
        allow_null_for_optional=allow_null_for_optional,
    )
    return ArgsModelBundle(model=model, parameters_schema=parameters_schema)
