# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

"""Audit vLLM compiled libraries for PyTorch stable ABI compliance."""

import fnmatch
import sys
from pathlib import Path

from torch_abi_audit import inspect_package
from torch_abi_audit.report import ExtensionReport, PackageReport

# Temporary allowlist of extensions not yet on the stable ABI.
# Shrink and remove over time.
ALLOWED_UNSTABLE_LIBRARIES: tuple[str, ...] = (
    # ROCm-only legacy extensions. Most ROCm ops have moved to
    # _C_stable_libtorch, but these two modules still bind libtorch directly.
    "_C.abi3.so",
    "_rocm_C.abi3.so",
)


def _relative_path(lib: ExtensionReport, package_root: Path) -> str:
    try:
        return lib.path.relative_to(package_root).as_posix()
    except ValueError:
        return lib.path.name


def _is_torch_unstable(lib: ExtensionReport) -> bool:
    return lib.error is None and lib.torch.uses_torch and not lib.torch.stable


def _matches_allowlist(rel_path: str, patterns: tuple[str, ...]) -> bool:
    return any(fnmatch.fnmatch(rel_path, pattern) for pattern in patterns)


def _iter_libs(report: PackageReport) -> tuple[ExtensionReport, ...]:
    return (*report.extensions, *report.bundled_libs)


def _collect_unstable(report: PackageReport) -> list[str]:
    return sorted(
        _relative_path(lib, report.root)
        for lib in _iter_libs(report)
        if _is_torch_unstable(lib)
    )


def _find_stale_allowlist_entries(
    report: PackageReport, patterns: tuple[str, ...]
) -> list[str]:
    """Allowlist patterns that match a built library which is no longer unstable."""
    stale: list[str] = []
    for pattern in patterns:
        for lib in _iter_libs(report):
            if lib.error is not None:
                continue
            if not fnmatch.fnmatch(_relative_path(lib, report.root), pattern):
                continue
            if not _is_torch_unstable(lib):
                stale.append(pattern)
                break
    return stale


def check_torch_abi(
    package: str = "vllm",
    patterns: tuple[str, ...] = ALLOWED_UNSTABLE_LIBRARIES,
) -> int:
    report = inspect_package(package)
    if report.error:
        print(f"error: failed to inspect {package!r}: {report.error}", file=sys.stderr)
        return 2

    unstable = _collect_unstable(report)
    unexpected = [
        rel_path for rel_path in unstable if not _matches_allowlist(rel_path, patterns)
    ]
    stale = _find_stale_allowlist_entries(report, patterns)

    if unexpected or stale:
        if unexpected:
            print(
                "Not allowed: torch-unstable libraries outside "
                f"ALLOWED_UNSTABLE_LIBRARIES: {', '.join(unexpected)}",
                file=sys.stderr,
            )
        if stale:
            print(
                "Not allowed: stale ALLOWED_UNSTABLE_LIBRARIES entries: "
                f"{', '.join(stale)}",
                file=sys.stderr,
            )
        return 1

    print("Torch stable ABI check passed.")
    return 0


if __name__ == "__main__":
    print(">>> Auditing vLLM extension modules for PyTorch stable ABI compliance")
    sys.exit(check_torch_abi())