"""
Annotation inspection and RST formatting.

Converts Python type annotations into reStructuredText cross-reference markup for Sphinx output.
Handles standard library types, typing constructs (Union, Optional, Literal, NewType, etc.),
forward references, type aliases, and pydata-sphinx-theme compatibility annotations.

This is a leaf module with no internal package imports to keep the dependency graph acyclic.
"""

from __future__ import annotations

import enum
import inspect
import re
import sys
import types
from typing import TYPE_CHECKING, Any, AnyStr, ForwardRef, NewType, TypeVar, Union, get_args, get_origin

from sphinx.util import rst
from sphinx.util.inspect import TypeAliasForwardRef

if TYPE_CHECKING:
    from collections.abc import Callable, Generator

    from sphinx.config import Config

from typing import TypeAliasType

_PYDATA_ANNOTS_TYPING = {
    "Any",
    "AnyStr",
    "Callable",
    "ClassVar",
    "Literal",
    "NoReturn",
    "Optional",
    "Tuple",
    *({"Union"} if sys.version_info < (3, 14) else set()),
}
# The types module documents these as ``py:data`` through 3.12 and as ``py:class`` from 3.13 on.
_PYDATA_ANNOTS_TYPES = (
    {
        *("AsyncGeneratorType", "BuiltinFunctionType", "BuiltinMethodType"),
        *("CellType", "ClassMethodDescriptorType", "CoroutineType"),
        *("EllipsisType", "FrameType", "FunctionType"),
        *("GeneratorType", "GetSetDescriptorType"),
        "LambdaType",
        *("MemberDescriptorType", "MethodDescriptorType", "MethodType", "MethodWrapperType"),
        *("NoneType", "NotImplementedType", "WrapperDescriptorType"),
    }
    if sys.version_info < (3, 13)
    else set()
)
_PYDATA_ANNOTATIONS = {
    *(("typing", n) for n in _PYDATA_ANNOTS_TYPING),
    *(("types", n) for n in _PYDATA_ANNOTS_TYPES),
}

_TYPES_DICT = {getattr(types, name): name for name in types.__all__}
_TYPES_DICT[types.FunctionType] = "FunctionType"

_UNESCAPE_RE = re.compile(
    r"""
    \\          # literal backslash
    ([^ ])      # followed by any non-space character (captured)
    """,
    re.VERBOSE,
)


class MyTypeAliasForwardRef(TypeAliasForwardRef):
    crossref: bool = False

    def __or__(self, value: Any) -> Any:  # ty: ignore[invalid-method-override]
        return Union[self, value]  # ruff:ignore[non-pep604-annotation-union]

    def __getitem__(self, item: Any) -> Any:
        # A generic alias stands in for the real one while resolving string annotations, so it has to
        # be subscriptable: evaluating ``"Pair[int]"`` would raise ``TypeError`` and drop the hint.
        return types.GenericAlias(self, item)  # ty: ignore[invalid-argument-type]  # any origin works at runtime


def format_annotation(annotation: Any, config: Config, *, short_literals: bool = False) -> str:
    """
    Format the annotation into reStructuredText cross-reference markup.

    The annotation is a tree (``dict[str, list[int]]`` nests three levels), so formatting it is a
    post-order walk. The walk is driven iteratively with an explicit stack rather than by recursion:
    a PEP 695 type alias may reference itself (``type T = int | list[T]``), which is a fully
    supported recursive type but would expand forever and exhaust the interpreter recursion limit if
    walked naively -- as may a generic one through its subscripted form (``type T[X] = X | list[T[X]]``).
    Each frame on the stack is one suspended ``_format_node`` generator; the generator
    yields the child annotations it needs formatted and is resumed with their rendered strings.
    ``active`` holds the aliases currently being expanded along the path from the root, so when an
    alias references itself it is rendered as a cross-reference to its own name -- the natural
    representation of a recursive type -- instead of being expanded again.
    """
    stack = [(_format_node(annotation, config, short_literals=short_literals), annotation)]
    active = {id(alias)} if (alias := _underlying_type_alias(annotation)) is not None else set()
    rendered: str | None = None
    while True:
        gen, node = stack[-1]
        try:
            child = next(gen) if rendered is None else gen.send(rendered)
        except StopIteration as done:
            stack.pop()
            if (alias := _underlying_type_alias(node)) is not None:
                active.discard(id(alias))
            rendered = done.value
            if not stack:
                return rendered
            continue
        if (alias := _underlying_type_alias(child)) is not None and id(alias) in active:
            rendered = _type_alias_crossref(alias, config) or _alias_name_reference(alias, config)
            continue
        stack.append((_format_node(child, config, short_literals=short_literals), child))
        if alias is not None:
            active.add(id(alias))
        rendered = None


def _format_node(  # ruff:ignore[complex-structure, too-many-return-statements, too-many-branches, too-many-statements, too-many-locals]
    annotation: Any,
    config: Config,
    *,
    short_literals: bool,
) -> Generator[Any, str, str]:
    """
    Format a single annotation node, delegating nested annotations to :func:`format_annotation`.

    Every nested annotation is ``yield``-ed and the driver sends back its formatted string, so the
    recursion lives in the driver's explicit stack rather than on the Python call stack.
    """
    typehints_formatter: Callable[..., str] | None = getattr(config, "typehints_formatter", None)
    if typehints_formatter is not None:
        formatted = typehints_formatter(annotation, config)
        if formatted is not None:
            return formatted

    if isinstance(annotation, ForwardRef):
        return annotation.__forward_arg__
    if annotation is None or annotation is type(None):
        return ":py:obj:`None`"
    if annotation is Ellipsis:
        return ":py:data:`...<Ellipsis>`"

    if isinstance(annotation, tuple):
        parts = []
        for item in annotation:
            parts.append((yield item))  # ruff:ignore[manual-list-comprehension]  # yield is illegal inside a comprehension
        if not parts:
            return "()"
        if len(parts) == 1:
            return f"({parts[0]}, )"
        return f"({', '.join(parts)})"

    if isinstance(annotation, TypeAliasForwardRef):
        fully_qualified: bool = getattr(config, "typehints_fully_qualified", False)
        prefix = "" if fully_qualified else "~"
        if (env := getattr(config, "_typehints_env", None)) is not None:
            py_domain = env.get_domain("py")
            module_prefix = getattr(config, "_typehints_module_prefix", "")
            for candidate in (f"{module_prefix}.{annotation.name}", annotation.name):
                if candidate in py_domain.objects and py_domain.objects[candidate].objtype == "type":
                    return f":py:type:`{prefix}{candidate}`"
        if isinstance(annotation, MyTypeAliasForwardRef) and annotation.crossref:
            return f":py:type:`{prefix}{annotation.name}`"
        return annotation.name

    if isinstance(annotation, TypeAliasType):
        if (crossref := _type_alias_crossref(annotation, config)) is not None:
            return crossref
        return (yield annotation.__value__)

    if isinstance(alias := get_origin(annotation), TypeAliasType | TypeAliasForwardRef):
        # A subscripted generic alias, e.g. ``NDArray[np.void]`` for ``type NDArray[ScalarT] = ...``.
        # It is a plain ``types.GenericAlias``, so without unwrapping it here the code below would
        # render it as its wrapper class' name instead of as the alias.
        args = get_args(annotation)
        if isinstance(alias, TypeAliasForwardRef):  # a documented alias, resolved from a string annotation
            rendered_alias = yield alias
        elif (crossref := _type_alias_crossref(alias, config)) is not None:
            rendered_alias = crossref
        elif _matches_type_params(alias, args):  # an undocumented alias: expand its value, with args substituted
            return (yield _substitute_type_params(alias, args))
        else:  # a wrong-arity subscript cannot expand, so name the alias instead of leaking its type params
            rendered_alias = _alias_name_reference(alias, config)
        parts = []
        for arg in args:
            parts.append((yield arg))  # yield is illegal inside a comprehension
        return f"{rendered_alias}\\ \\[{', '.join(parts)}]"

    try:
        module = get_annotation_module(annotation)
        class_name = get_annotation_class_name(annotation, module)
        args = get_annotation_args(annotation, module, class_name)
    except ValueError:
        return str(annotation).strip("'")

    module = _fixup_module_name(config, module)
    internal_path = f"{module}.{class_name}"
    if (mapping := getattr(config, "_intersphinx_type_mapping", None)) and (mapped := mapping.get(internal_path)):
        module, _, class_name = mapped.rpartition(".")
    full_name = f"{module}.{class_name}" if module != "builtins" else class_name
    fully_qualified = getattr(config, "typehints_fully_qualified", False)
    prefix = "" if fully_qualified or full_name == class_name else "~"
    role = "data" if (module, class_name) in _PYDATA_ANNOTATIONS else "class"
    args_format = "\\[{}]"
    formatted_args: str | None = ""

    always_use_bars_union: bool = getattr(config, "always_use_bars_union", True)
    is_bars_union = (
        (sys.version_info >= (3, 14) and full_name == "typing.Union")
        or full_name == "types.UnionType"
        or (always_use_bars_union and type(annotation).__qualname__ == "_UnionGenericAlias")
    )
    if is_bars_union:
        full_name = ""

    if full_name == "typing.NewType":
        newtype_module = _fixup_module_name(config, getattr(annotation, "__module__", ""))
        newtype_name = annotation.__name__
        newtype_qualified = f"{newtype_module}.{newtype_name}" if newtype_module else newtype_name
        newtype_prefix = "" if fully_qualified or not newtype_module else "~"
        supertype = yield annotation.__supertype__
        return f":py:class:`{newtype_prefix}{newtype_qualified}` ({supertype})"
    if full_name == "typing.Annotated":
        return (yield annotation.__origin__)
    if full_name in {"typing.TypeVar", "typing.ParamSpec"}:
        params = {k: getattr(annotation, f"__{k}__") for k in ("bound", "covariant", "contravariant")}
        params = {k: v for k, v in params.items() if v}
        if "bound" in params:
            bound = yield params["bound"]
            params["bound"] = f" {bound}"
        args_format = f"\\(``{annotation.__name__}``{', {}' if args else ''}"
        if params:
            args_format += "".join(f", {k}={v}" for k, v in params.items())
        args_format += ")"
        formatted_args = None if args else args_format
    elif full_name == "typing.Optional":  # pragma: <3.14 cover
        args = tuple(x for x in args if x is not type(None))
    elif full_name in {"typing.Union", "types.UnionType"} and type(None) in args:  # pragma: <3.14 cover
        if len(args) == 2:  # ruff:ignore[magic-value-comparison]
            full_name = "typing.Optional"
            role = "data"
            args = tuple(x for x in args if x is not type(None))
        else:
            simplify_optional_unions: bool = getattr(config, "simplify_optional_unions", True)
            if not simplify_optional_unions:
                full_name = "typing.Optional"
                role = "data"
                args_format = f"\\[:py:data:`{prefix}typing.Union`\\[{{}}]]"
                args = tuple(x for x in args if x is not type(None))
    elif full_name in {"typing.Callable", "collections.abc.Callable"} and args and args[0] is not ...:
        fmt = []
        for arg in args:
            fmt.append((yield arg))  # ruff:ignore[manual-list-comprehension]  # yield is illegal inside a comprehension
        formatted_args = f"\\[\\[{', '.join(fmt[:-1])}], {fmt[-1]}]"
    elif full_name == "typing.Literal":
        literal_parts = [_format_literal_arg(arg, config) for arg in args]
        if short_literals:
            return f"\\{' | '.join(literal_parts)}"
        formatted_args = f"\\[{', '.join(literal_parts)}]"
    elif is_bars_union:
        if not args:
            return f":py:{'class' if sys.version_info >= (3, 14) else 'data'}:`{prefix}typing.Union`"
        union_parts = []
        for arg in args:
            union_parts.append((yield arg))  # ruff:ignore[manual-list-comprehension]  # yield is illegal inside a comprehension
        return " | ".join(union_parts)

    if args and not formatted_args:
        fmt = []
        for arg in args:
            fmt.append((yield arg))
        formatted_args = args_format.format(", ".join(fmt))

    escape = "\\ " if formatted_args else ""
    # B901: returning a value from a generator is intentional here — the driver reads it via StopIteration
    return f":py:{role}:`{prefix}{full_name}`{escape}{formatted_args}"  # ruff:ignore[return-in-generator]


def get_annotation_module(annotation: Any) -> str:
    if annotation is None:
        return "builtins"
    if _get_types_type(annotation) is not None:
        return "types"
    is_new_type = isinstance(annotation, NewType)
    if (
        is_new_type
        or isinstance(annotation, TypeVar)
        or type(annotation).__name__ in {"ParamSpec", "ParamSpecArgs", "ParamSpecKwargs"}
    ):
        return "typing"
    if hasattr(annotation, "__module__"):
        return annotation.__module__
    msg = f"Cannot determine the module of {annotation}"
    raise ValueError(msg)


def get_annotation_class_name(annotation: Any, module: str) -> str:  # ruff:ignore[complex-structure, too-many-return-statements]
    if annotation is None:
        return "None"
    if annotation is AnyStr:
        return "AnyStr"
    val = _get_types_type(annotation)
    if val is not None:
        return val
    if _is_newtype(annotation):
        return "NewType"

    if getattr(annotation, "__qualname__", None):
        return annotation.__qualname__
    if getattr(annotation, "_name", None):  # pragma: <3.14 cover
        return annotation._name  # ruff:ignore[private-member-access]
    if module in {"typing", "typing_extensions"} and isinstance(
        getattr(annotation, "name", None), str
    ):  # pragma: <3.14 cover
        return annotation.name

    origin = getattr(annotation, "__origin__", None)
    if origin:
        if getattr(origin, "__qualname__", None):  # pragma: <3.14 cover
            return origin.__qualname__
        if getattr(origin, "_name", None):  # pragma: <3.14 cover
            return origin._name  # ruff:ignore[private-member-access]

    annotation_cls = annotation if inspect.isclass(annotation) else type(annotation)
    return annotation_cls.__qualname__.lstrip("_")


def get_annotation_args(annotation: Any, module: str, class_name: str) -> tuple[Any, ...]:
    try:
        original = getattr(sys.modules[module], class_name)
    except (KeyError, AttributeError):
        pass
    else:
        if annotation is original:
            return ()

    if class_name == "TypeVar" and hasattr(annotation, "__constraints__"):
        return annotation.__constraints__
    if class_name == "NewType" and hasattr(annotation, "__supertype__"):
        return (annotation.__supertype__,)
    if class_name == "Generic":
        return annotation.__parameters__
    result = getattr(annotation, "__args__", ())
    return () if len(result) == 1 and result[0] == () else result  # type: ignore[misc]


def _fixup_module_name(config: Config, module: str) -> str:
    if getattr(config, "typehints_fixup_module_name", None):
        module = config.typehints_fixup_module_name(module)
    if module == "typing_extensions":  # pragma: <3.14 cover
        module = "typing"
    if module == "_io":
        module = "io"
    return module


def _format_literal_arg(arg: Any, config: Config) -> str:
    if isinstance(arg, enum.Enum):
        enum_cls = type(arg)
        module = _fixup_module_name(config, enum_cls.__module__)
        fully_qualified = getattr(config, "typehints_fully_qualified", False)
        qualified = f"{module}.{enum_cls.__qualname__}.{arg.name}" if module else f"{enum_cls.__qualname__}.{arg.name}"
        prefix = "" if fully_qualified or not module else "~"
        return f":py:attr:`{prefix}{qualified}`"
    return f"``{arg!r}``"


def _get_types_type(obj: Any) -> str | None:
    try:
        return _TYPES_DICT.get(obj)
    except Exception:  # ruff:ignore[blind-except]
        return None


def _is_newtype(annotation: Any) -> bool:
    return isinstance(annotation, NewType)


def _underlying_type_alias(annotation: Any) -> TypeAliasType | None:
    """Return the PEP 695 alias an annotation is, or the one a subscripted alias parametrizes."""
    if isinstance(annotation, TypeAliasType):
        return annotation
    if isinstance(origin := get_origin(annotation), TypeAliasType):
        return origin
    return None


def _matches_type_params(alias: TypeAliasType, args: tuple[Any, ...]) -> bool:
    """Whether ``args`` can be substituted: the runtime accepts a wrong-arity subscript, so check it here."""
    params = alias.__type_params__
    if any(not isinstance(param, TypeVar) for param in params):  # a TypeVarTuple or ParamSpec absorbs any count
        return True
    return len(params) == len(args)


def _substitute_type_params(alias: TypeAliasType, args: tuple[Any, ...]) -> Any:
    """Substitute ``args`` into an undocumented generic alias' value, so it expands to concrete types."""
    if isinstance(value := alias.__value__, TypeVar):
        # ``type Ident[X] = X`` has a bare type param as its value: it is the substitution, not a container
        return dict(zip(alias.__type_params__, args, strict=False)).get(value, value)
    try:
        return value[args]
    except TypeError:  # the value ignores its type params, e.g. ``type Count[X] = int``
        return value


def _alias_name_reference(alias: TypeAliasType, config: Config) -> str:
    """
    Render an alias as a cross-reference to its own name, without expanding its value.

    A cross-module alias may live in a module Sphinx has not read yet, leaving it out of the py domain. Its
    canonical qualified name resolves in the later xref phase and renders the same short form when the alias
    is undocumented.
    """
    prefix = "" if getattr(config, "typehints_fully_qualified", False) else "~"
    return f":py:type:`{prefix}{_get_canonical_type_alias_name(alias) or alias.__name__}`"


def _type_alias_crossref(annotation: TypeAliasType, config: Config) -> str | None:
    """
    Render a PEP 695 type alias as a reStructuredText cross-reference, or ``None`` when it is not a known target.

    Look the alias up in the py domain under the names it could be documented as: walk up the current module
    prefix (so a ``T`` documented in the module being built wins), then its canonical module (so an alias
    imported from elsewhere still resolves), then its bare name. An alias owned by a different top-level package
    falls back to an intersphinx-style reference. Returning ``None`` tells the caller to expand the alias value.
    """
    env = getattr(config, "_typehints_env", None)
    if env is None:
        return None
    fully_qualified: bool = getattr(config, "typehints_fully_qualified", False)
    prefix = "" if fully_qualified else "~"
    py_domain = env.get_domain("py")
    module_prefix = getattr(config, "_typehints_module_prefix", "")
    prefix_parts = module_prefix.split(".") if module_prefix else []
    canonical = _get_canonical_type_alias_name(annotation)
    candidates = [".".join((*prefix_parts[:n], annotation.__name__)) for n in range(len(prefix_parts), 0, -1)]
    if canonical:
        candidates.append(canonical)
    candidates.append(annotation.__name__)
    for candidate in candidates:
        if (obj := py_domain.objects.get(candidate)) is not None and obj.objtype == "type":
            return f":py:type:`{prefix}{candidate}`"
    current_top = prefix_parts[0] if prefix_parts else ""
    if canonical and canonical.split(".")[0] != current_top:
        full_name = _fixup_module_name(config, canonical.rpartition(".")[0]) + "." + annotation.__name__
        return f":py:obj:`{prefix}{full_name}`"
    return None


def _get_canonical_type_alias_name(annotation: TypeAliasType) -> str:
    """
    Get canonical public qualified name for a TypeAliasType.

    For types defined in private modules (e.g. ``numpy._typing.ArrayLike``),
    search ``sys.modules`` for a public re-export
    (e.g. ``numpy.typing.ArrayLike``).
    """
    module = getattr(annotation, "__module__", "") or ""
    name = getattr(annotation, "__name__", "") or ""
    if not module or not name:
        return ""
    if not any(part.startswith("_") for part in module.split(".")):
        return f"{module}.{name}"
    top_pkg = module.split(".")[0]
    for mod_name in sorted(sys.modules):
        if not mod_name.startswith(top_pkg):
            continue
        mod = sys.modules[mod_name]
        if not isinstance(mod, types.ModuleType):
            continue
        if any(part.startswith("_") for part in mod_name.split(".")):
            continue
        if getattr(mod, name, None) is annotation:
            return f"{mod_name}.{name}"
    return f"{module}.{name}"


def unescape(escaped: str) -> str:
    escaped = escaped.replace("\x00", "")
    return _UNESCAPE_RE.sub(r"\1", escaped)


def add_type_css_class(type_rst: str) -> str:
    return f":sphinx_autodoc_typehints_type:`{rst.escape(type_rst)}`"
