"""Schema generation for OpenAI compatible APIs.

This module provides functionality to generate JSON schemas for classes and types
and to validate instances against these schemas.
"""

from dataclasses import is_dataclass
from enum import Enum
from types import FunctionType, UnionType
from typing import Callable, Union, get_args, get_origin, get_type_hints

from pydantic import TypeAdapter


_SCHEMA_OBJECT_WRAPPER = "schema_object_wrapper"
_SCHEMA_DICT_WRAPPER = "schema_dict_wrapper"


def _type_to_schema(ty: type, title: str = "", desc: str = "") -> dict:

    title = title.replace("_", " ").title()
    context = ({"title": title} if title else {}) | (
        {"description": desc} if desc else {}
    )

    semstr: str = ty._jac_semstr if hasattr(ty, "_jac_semstr") else ""
    semstr = semstr or (ty.__doc__ if hasattr(ty, "__doc__") else "") or ""  # type: ignore

    semstr_inner: dict[str, str] = (
        ty._jac_semstr_inner if hasattr(ty, "_jac_semstr_inner") else {}
    )

    # Raise on unsupported types
    if ty in (list, dict, set, tuple):
        raise ValueError(
            f"Untyped {ty.__name__} is not supported for schema generation. "
            f"Use {ty.__name__}[T, ...] instead."
        )

    # Handle primitive types
    if ty is type(None):
        return {"type": "null"} | context
    if ty is bool:
        return {"type": "boolean"} | context
    if ty is int:
        return {"type": "integer"} | context
    if ty is float:
        return {"type": "number"} | context
    if ty is str:
        return {"type": "string"} | context

    # Handle Union
    if get_origin(ty) in (Union, UnionType):
        args = get_args(ty)
        return {
            "anyOf": [_type_to_schema(arg) for arg in args],
            "title": title,
        } | context

    # Handle annotated list
    if get_origin(ty) is list:
        item_type: type = get_args(ty)[0]
        return {
            "type": "array",
            "items": _type_to_schema(item_type),
        } | context

    # Handle annotated tuple/set
    if get_origin(ty) in (tuple, set):
        origin = get_origin(ty).__name__  # type: ignore
        args = get_args(ty)
        if len(args) == 2 and args[1] is Ellipsis:
            item_type = args[0]
            return {
                "type": "array",
                "items": _type_to_schema(item_type),
            } | context
        raise ValueError(
            f"Unsupported {origin} type for schema generation: {ty}. "
            f"Only {origin} of the form {origin}[T, ...] are supported."
        )

    # Handle annotated dictionaries
    if get_origin(ty) is dict:
        return _convert_dict_to_schema(ty) | context

    # Handle dataclass
    if is_dataclass(ty):
        fields: dict[str, type] = {
            name: type
            for name, type in get_type_hints(ty).items()
            if not name.startswith("_")
        }
        properties = {
            name: _type_to_schema(type, name, semstr_inner.get(name, ""))  # type: ignore
            for name, type in fields.items()
        }
        return {
            "title": title or ty.__name__,
            "description": semstr,
            "type": "object",
            "properties": properties,
            "required": list(fields.keys()),
            "additionalProperties": False,
        }

    # Handle enums
    if isinstance(ty, type) and issubclass(ty, Enum):
        enum_type = None
        enum_values = []
        for member in ty.__members__.values():
            enum_values.append(member.value)
            if enum_type is None:
                enum_type = type(member.value)
            elif type(member.value) is not enum_type:
                raise ValueError(
                    f"Enum {ty.__name__} has mixed types. Not supported for schema generation."
                )
            enum_type = enum_type or int
        enum_desc = f"\nThe value *should* be one in this list: {enum_values}"
        if enum_type not in (int, str):
            raise ValueError(
                f"Enum {ty.__name__} has unsupported type {enum_type}. "
                "Only int and str enums are supported for schema generation."
            )
        return {
            "description": semstr + enum_desc,
            "type": "integer" if enum_type is int else "string",
        }

    # Handle functions
    if isinstance(ty, FunctionType):
        hints = get_type_hints(ty)
        hints.pop("return", None)
        params = {
            name: _type_to_schema(type, name, semstr_inner.get(name, ""))
            for name, type in hints.items()
        }
        return {
            "title": title or ty.__name__,
            "type": "function",
            "description": semstr,
            "properties": params,
            "required": list(params.keys()),
            "additionalProperties": False,
        }

    raise ValueError(
        f"Unsupported type for schema generation: {ty}. "
        "Only primitive types, dataclasses, and Union types are supported."
    )


def _name_of_type(ty: type) -> str:
    if get_origin(ty) in (Union, UnionType):
        names = [_name_of_type(arg) for arg in get_args(ty)]
        return "_or_".join(names)
    if hasattr(ty, "__name__"):
        return ty.__name__
    return "type"


def _convert_dict_to_schema(ty_dict: type) -> dict:
    """Convert a dictionary type to a schema."""
    if get_origin(ty_dict) is not dict:
        raise ValueError(f"Expected a dictionary type, got {ty_dict}.")
    key_type, value_type = get_args(ty_dict)
    return {
        "type": "object",
        "title": _SCHEMA_DICT_WRAPPER,
        "properties": {
            _SCHEMA_DICT_WRAPPER: {
                "type": "array",
                "items": {
                    "type": "object",
                    "properties": {
                        "key": _type_to_schema(key_type),
                        "value": _type_to_schema(value_type),
                    },
                    "required": ["key", "value"],
                    "additionalProperties": False,
                },
            }
        },
        "additionalProperties": False,
        "required": [_SCHEMA_DICT_WRAPPER],
    }


def _decode_dict(json_obj: dict) -> dict:
    """Decode a JSON dictionary to a Python dictionary."""
    if not isinstance(json_obj, dict):
        return json_obj
    if _SCHEMA_DICT_WRAPPER in json_obj:
        items = json_obj[_SCHEMA_DICT_WRAPPER]
        return {item["key"]: _decode_dict(item["value"]) for item in items}
    return {key: _decode_dict(value) for key, value in json_obj.items()}


def _wrap_to_object(schema: dict[str, object]) -> dict[str, object]:
    """Wrap the schema in an object with a type."""
    if "type" in schema and schema["type"] == "object":
        return schema
    return {
        "type": "object",
        "title": _SCHEMA_OBJECT_WRAPPER,
        "properties": {
            _SCHEMA_OBJECT_WRAPPER: schema,
        },
        "required": [_SCHEMA_OBJECT_WRAPPER],
        "additionalProperties": False,
    }


def _unwrap_from_object(json_obj: dict) -> dict:
    """Unwrap the schema from an object with a type."""
    if _SCHEMA_OBJECT_WRAPPER in json_obj:
        return json_obj[_SCHEMA_OBJECT_WRAPPER]
    return json_obj


def type_to_schema(resp_type: type) -> dict[str, object]:
    """Return the JSON schema for the response type."""
    type_name = _name_of_type(resp_type)
    schema = _type_to_schema(resp_type, type_name)
    schema = _wrap_to_object(schema)
    return {
        "type": "json_schema",
        "json_schema": {
            "name": type_name,
            "schema": schema,
            "strict": True,
        },
    }


def tool_to_schema(
    func: Callable, description: str, params_desc: dict[str, str]
) -> dict[str, object]:
    """Return the JSON schema for the tool type."""
    schema = _type_to_schema(func)  # type: ignore
    properties: dict[str, object] = schema.get("properties", {})  # type: ignore
    required: list[str] = schema.get("required", [])  # type: ignore
    for param_name, param_info in properties.items():
        param_info["description"] = params_desc.get(param_name, "")  # type: ignore
    return {
        "type": "function",
        "function": {
            "name": func.__name__,
            "description": description,
            "parameters": {
                "type": "object",
                "properties": properties,
                "required": required,
                "additionalProperties": False,
            },
        },
    }


def json_to_instance(json_obj: dict, ty: type) -> object:
    """Convert a JSON dictionary to an instance of the given type."""
    json_obj = _unwrap_from_object(json_obj)
    json_obj = _decode_dict(json_obj)
    return TypeAdapter(ty).validate_python(json_obj)
