from __future__ import annotations

# standard
import sys
import json
import typing
import inspect
import types as _builtin_types

# third party
# custom
from sunwaee.modules.gen.engine.types import Tool


def _is_union(origin: typing.Any) -> bool:
    """Check if is union."""

    if origin is typing.Union:
        return True
    if sys.version_info >= (3, 10):
        return origin is _builtin_types.UnionType
    return False  # pragma: no cover


def _unwrap_optional(annotation: typing.Any) -> tuple[typing.Any, bool]:
    """Strip `None` from a union. Returns (inner_type, is_optional)."""

    origin = typing.get_origin(annotation)
    if _is_union(origin):
        args = [a for a in typing.get_args(annotation) if a is not type(None)]
        if len(args) == 1:
            return args[0], True
    return annotation, False


def _py_to_jsonschema(annotation: typing.Any) -> dict[str, typing.Any]:
    """Transform python prop to json."""

    origin = typing.get_origin(annotation)
    args = typing.get_args(annotation)

    # Annotated[T, "description", ...]
    if origin is typing.Annotated:
        schema = _py_to_jsonschema(args[0])
        for meta in args[1:]:
            if isinstance(meta, str):
                schema["description"] = meta
        return schema

    # Union / Optional --> unwrap None, recurse on inner type
    if _is_union(origin):
        inner, _ = _unwrap_optional(annotation)
        return _py_to_jsonschema(inner)

    # Literal["a", "b"]
    if origin is typing.Literal:
        return {"type": "string", "enum": list(args)}

    # list[T]
    if origin is list:
        item_schema = _py_to_jsonschema(args[0]) if args else {"type": "string"}
        return {"type": "array", "items": item_schema}

    _PRIMITIVES = {str: "string", int: "integer", float: "number", bool: "boolean"}
    if annotation in _PRIMITIVES:
        return {"type": _PRIMITIVES[annotation]}

    return {"type": "string"}


def _is_optional_annotation(annotation: typing.Any) -> bool:
    """Return True if the annotation allows None (Optional / X | None)."""

    # unwrap Annotated first
    if typing.get_origin(annotation) is typing.Annotated:
        annotation = typing.get_args(annotation)[0]
    _, is_optional = _unwrap_optional(annotation)
    return is_optional


###  DECORATOR


def tool(description: str) -> typing.Callable:
    """Decorator that attaches a Tool definition to the function as `fn._tool`."""

    def decorator(fn: typing.Callable) -> typing.Callable:
        hints = typing.get_type_hints(fn, include_extras=True)
        sig = inspect.signature(fn)

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

        for name, param in sig.parameters.items():
            annotation = hints.get(name, str)
            properties[name] = _py_to_jsonschema(annotation)

            has_default = param.default is not inspect.Parameter.empty
            if not has_default and not _is_optional_annotation(annotation):
                required.append(name)

        fn._tool = Tool(
            name=fn.__name__,
            description=description,
            parameters={
                "type": "object",
                "properties": properties,
                "required": required,
            },
            fn=fn,
        )
        return fn

    return decorator


### HELPERS


def ok(data: typing.Any, git_diff: str | None = None) -> str:
    """Successful tool call output."""

    result: dict[str, typing.Any] = {"ok": True, "data": data}
    if git_diff:
        result["git_diff"] = git_diff
    return json.dumps(result, default=str)


def err(message: str) -> str:
    """Error tool call output."""

    return json.dumps({"ok": False, "error": message})
