# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

"""Internal support for error-enum explanations.

``cuda_core`` keeps frozen 13.1.1 fallback tables for older ``cuda-bindings``
releases. Driver/runtime error enums carry usable ``__doc__`` text starting in
the 12.x backport line at ``cuda-bindings`` 12.9.6, and in the mainline 13.x
series at ``cuda-bindings`` 13.2.0. This module decides which source to use
and normalizes generated docstrings so user-facing ``CUDAError`` messages stay
presentable.

The cleanup rules here were derived while validating generated enum docstrings
in PR #1805. Keep them narrow and remove them when codegen quirks or fallback
support are no longer needed.
"""

from __future__ import annotations

import importlib.metadata
import re
from collections.abc import Callable
from typing import Any

_MIN_12X_BINDING_VERSION_FOR_ENUM_DOCSTRINGS = (12, 9, 6)
_MIN_13X_BINDING_VERSION_FOR_ENUM_DOCSTRINGS = (13, 2, 0)
_RST_INLINE_ROLE_RE = re.compile(r":(?:[a-z]+:)?[a-z]+:`([^`]+)`")
_WORDWRAP_HYPHEN_AFTER_RE = re.compile(r"(?<=[0-9A-Za-z_])- (?=[0-9A-Za-z_])")
_WORDWRAP_HYPHEN_BEFORE_RE = re.compile(r"(?<=[0-9A-Za-z_]) -(?=[0-9A-Za-z_])")
_ExplanationTable = dict[int, str | tuple[str, ...]]
_ExplanationTableLoader = Callable[[], _ExplanationTable]


def _parse_version_triple(version_str: str) -> tuple[int, int, int]:
    """Parse a PEP 440 version string into a (major, minor, patch) triple.

    Strips local-version identifiers and handles pre-release suffixes such as
    ``0b1`` or ``0rc1`` by extracting only the leading integer from each
    release segment.
    """
    parts = version_str.partition("+")[0].split(".")[:3]
    ints = ([int(m.group(1)) if (m := re.match(r"(\d+)", v)) else 0 for v in parts] + [0, 0, 0])[:3]
    return (ints[0], ints[1], ints[2])


# ``version.pyx`` cannot be reused here (circular import via ``cuda_utils``).
def _binding_version() -> tuple[int, int, int]:
    """Return the installed ``cuda-bindings`` version, or a conservative old value."""
    try:
        version = importlib.metadata.version("cuda-bindings")
    except importlib.metadata.PackageNotFoundError:
        return (0, 0, 0)  # For very old versions of cuda-python
    return _parse_version_triple(version)


def _binding_version_has_usable_enum_docstrings(version: tuple[int, int, int]) -> bool:
    """Whether released bindings are known to carry usable error-enum ``__doc__`` text."""
    return (
        _MIN_12X_BINDING_VERSION_FOR_ENUM_DOCSTRINGS <= version < (13, 0, 0)
        or version >= _MIN_13X_BINDING_VERSION_FOR_ENUM_DOCSTRINGS
    )


def _fix_hyphenation_wordwrap_spacing(s: str) -> str:
    """Remove spaces around hyphens introduced by line wrapping in generated ``__doc__`` text.

    This targets asymmetric wrap artifacts such as ``non- linear`` or
    ``GPU- Direct`` while leaving intentional ``a - b`` separators alone.
    """
    prev = None
    while prev != s:
        prev = s
        s = _WORDWRAP_HYPHEN_AFTER_RE.sub("-", s)
        s = _WORDWRAP_HYPHEN_BEFORE_RE.sub("-", s)
    return s


def clean_enum_member_docstring(doc: str | None) -> str | None:
    """Turn an enum member ``__doc__`` into plain text.

    The generated enum docstrings are already close to user-facing prose, but
    they may contain Sphinx inline roles, line wrapping, or a small known
    codegen defect. Normalize only those differences so the text is suitable
    for error messages.
    """
    if doc is None:
        return None
    s = doc
    # Known codegen bug on cudaErrorIncompatibleDriverContext. Remove once fixed
    # in cuda-bindings code generation.
    s = s.replace("\n:py:obj:`~.Interactions`", ' "Interactions ')
    # Drop a leading "~." or "." after removing the surrounding RST inline role.
    s = _RST_INLINE_ROLE_RE.sub(lambda m: re.sub(r"^~?\.", "", m.group(1)), s)
    # Strip simple bold emphasis markers.
    s = re.sub(r"\*\*([^*]+)\*\*", r"\1", s)
    # Strip simple italic emphasis markers.
    s = re.sub(r"\*([^*]+)\*", r"\1", s)
    # Collapse wrapped lines and repeated spaces.
    s = re.sub(r"\s+", " ", s).strip()
    s = _fix_hyphenation_wordwrap_spacing(s)
    return s


class DocstringBackedExplanations:
    """Compatibility shim exposing enum-member ``__doc__`` text via ``dict.get``.

    Keeps the existing ``.get(int(error))`` lookup shape used by ``cuda_utils.pyx``.
    """

    __slots__ = ("_enum_type",)

    def __init__(self, enum_type: Any) -> None:
        self._enum_type = enum_type

    def get(self, code: int, default: str | None = None) -> str | None:
        try:
            member = self._enum_type(code)
        except ValueError:
            return default

        raw_doc = member.__doc__
        if raw_doc is None:
            return default

        return clean_enum_member_docstring(raw_doc)


def get_best_available_explanations(
    enum_type: Any,
    fallback: _ExplanationTable | _ExplanationTableLoader,
) -> DocstringBackedExplanations | _ExplanationTable:
    """Pick one explanation source per bindings version.

    Use enum-member ``__doc__`` only for bindings versions known to expose
    usable per-member text (12.9.6+ in the 12.x backport line, 13.2.0+ in the
    13.x mainline). Otherwise keep using the frozen 13.1.1 fallback tables.
    """
    if not _binding_version_has_usable_enum_docstrings(_binding_version()):
        if callable(fallback):
            return fallback()
        return fallback
    return DocstringBackedExplanations(enum_type)
