#!/usr/bin/env python3
"""
CPU-runnable unit tests for CUDA compute capability -> sm_XX mapping logic.

These tests replicate the C++ compute_cap_to_sm() and
identify_cuda_arch_from_name() functions so the mapping can be validated
without NVIDIA hardware or a running server.

Run with: python -m pytest test/test_cuda_arch_mapping.py
      or: python test/test_cuda_arch_mapping.py
"""

import unittest

# ---------------------------------------------------------------------------
# Python replica of system_info.cpp::compute_cap_to_sm()
# ---------------------------------------------------------------------------

CUDA_SUPPORTED_ARCHS = {
    "sm_75",  # Turing
    "sm_80",  # Ampere DC
    "sm_86",  # Ampere
    "sm_89",  # Ada Lovelace
    "sm_90",  # Hopper
    "sm_100",  # Blackwell DC
    "sm_120",  # Blackwell consumer
    "sm_121",  # Blackwell GB10 / Thor SoC
}


def compute_cap_to_sm(compute_cap: str) -> str:
    """
    Convert a "MAJOR.MINOR" compute-capability string from nvidia-smi to the
    sm_XX token used in llamacpp-cuda release asset filenames.
    Returns "" if the value cannot be parsed.
    """
    dot = compute_cap.find(".")
    if dot == -1:
        return ""
    major, minor = compute_cap[:dot], compute_cap[dot + 1 :]
    if not major or not minor:
        return ""
    try:
        return f"sm_{int(major) * 10 + int(minor)}"
    except ValueError:
        return ""


# ---------------------------------------------------------------------------
# Python replica of system_info.cpp::identify_cuda_arch_from_name() (compact)
# ---------------------------------------------------------------------------

# Data-center Blackwell (B100/B200) is compute capability 10.0 (sm_100), not
# 12.0 (sm_120) like the consumer/workstation Blackwell parts. These are matched
# before the generic "blackwell" keyword below so a name like
# "NVIDIA B200 (Blackwell)" still resolves to sm_100. Mirrors the C++ pre-check.
_DC_BLACKWELL = ["b100", "b200"]

_NAME_ARCH_TABLE = [
    # GB10 Thor SoC (sm_121)
    (["gb10"], "sm_121"),
    # Blackwell consumer/workstation (sm_120). The generic "blackwell" keyword
    # covers current and future RTX 50 / RTX PRO Blackwell GPUs; explicit model
    # numbers cover names that do not include "Blackwell".
    (
        [
            "blackwell",
            "rtx 50",
            "rtx50",
            "5090",
            "5080",
            "5070",
            "5060",
            "rtx pro 6000",
            "rtx pro 5000",
            "rtx pro 4500",
            "rtx pro 4000",
            "rtx pro 3500",
            "rtx pro 3000",
            "rtx pro 2000",
            "rtx pro 1000",
        ],
        "sm_120",
    ),
    # Hopper (sm_90)
    (["h100", "h200"], "sm_90"),
    # Ada Lovelace (sm_89)
    (["rtx 40", "rtx40", "4090", "4080", "4070", "4060", "l40", " l4"], "sm_89"),
    # Ampere DC (sm_80)
    (["a100"], "sm_80"),
    # Ampere consumer/pro (sm_86)
    (
        [
            "rtx 30",
            "rtx30",
            "3090",
            "3080",
            "3070",
            "3060",
            "3050",
            "a40",
            "a30",
            "a10",
            "a6000",
            "a5000",
            "a4000",
            "a2000",
        ],
        "sm_86",
    ),
    # Turing (sm_75)
    (
        [
            "rtx 20",
            "rtx20",
            "2080",
            "2070",
            "2060",
            "gtx 16",
            "gtx16",
            "1660",
            "1650",
            "titan rtx",
            "quadro rtx",
            " t4",
        ],
        "sm_75",
    ),
]


def identify_cuda_arch_from_name(device_name: str) -> str:
    """Compact fallback: infer sm_XX from GPU marketing name."""
    name = device_name.lower()
    # Require at least one NVIDIA identifier
    nvidia_ids = [
        "nvidia",
        "geforce",
        "rtx",
        "gtx",
        "quadro",
        "tesla",
        "titan",
        "a100",
        "a40",
        "a30",
        "a10",
        "h100",
        "h200",
        "b100",
        "b200",
        "l40",
        "gb10",
    ]
    if not any(kw in name for kw in nvidia_ids):
        return ""
    # Data-center Blackwell first, before the generic "blackwell" keyword.
    if any(kw in name for kw in _DC_BLACKWELL):
        return "sm_100"
    for keywords, sm in _NAME_ARCH_TABLE:
        if any(kw in name for kw in keywords):
            return sm
    return ""


# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------


class TestComputeCapToSm(unittest.TestCase):
    def test_valid_mappings(self):
        cases = [
            ("7.5", "sm_75"),
            ("8.0", "sm_80"),
            ("8.6", "sm_86"),
            ("8.9", "sm_89"),
            ("9.0", "sm_90"),
            ("10.0", "sm_100"),
            ("12.0", "sm_120"),
            ("12.1", "sm_121"),
        ]
        for cap, expected in cases:
            with self.subTest(cap=cap):
                self.assertEqual(compute_cap_to_sm(cap), expected)

    def test_unsupported_but_parseable(self):
        # Should still parse even if not in CUDA_SUPPORTED_ARCHS
        self.assertEqual(compute_cap_to_sm("6.1"), "sm_61")
        self.assertEqual(compute_cap_to_sm("5.0"), "sm_50")

    def test_invalid_inputs(self):
        for bad in ["", "86", "8", ".", "abc", "8.x"]:
            with self.subTest(bad=bad):
                self.assertEqual(compute_cap_to_sm(bad), "")

    def test_supported_arch_membership(self):
        for cap in ["7.5", "8.0", "8.6", "8.9", "9.0", "10.0", "12.0", "12.1"]:
            sm = compute_cap_to_sm(cap)
            self.assertIn(
                sm, CUDA_SUPPORTED_ARCHS, f"{cap} -> {sm} not in supported set"
            )

    def test_unsupported_arch_not_in_set(self):
        for cap in ["6.1", "7.0", "5.0"]:
            sm = compute_cap_to_sm(cap)
            self.assertNotIn(
                sm, CUDA_SUPPORTED_ARCHS, f"{cap} -> {sm} unexpectedly in supported set"
            )


class TestIdentifyCudaArchFromName(unittest.TestCase):
    def test_known_consumer_gpus(self):
        cases = [
            ("NVIDIA GeForce RTX 3080", "sm_86"),
            ("NVIDIA GeForce RTX 4090", "sm_89"),
            ("NVIDIA GeForce RTX 2080 Ti", "sm_75"),
            ("NVIDIA GeForce RTX 5090", "sm_120"),
            ("NVIDIA GTX 1660 Super", "sm_75"),
            ("NVIDIA H100 PCIe", "sm_90"),
            ("NVIDIA A100-SXM4-80GB", "sm_80"),
            ("NVIDIA A10", "sm_86"),
            ("NVIDIA A40", "sm_86"),
            ("NVIDIA GB10", "sm_121"),
            ("GB10", "sm_121"),
            # RTX PRO Blackwell (sm_120) — both the generic "Blackwell" keyword
            # and the explicit RTX PRO model numbers must resolve to sm_120.
            ("NVIDIA RTX PRO 3000 Blackwell Generation Laptop GPU", "sm_120"),
            ("NVIDIA RTX PRO 6000 Blackwell", "sm_120"),
            ("NVIDIA RTX PRO 1000", "sm_120"),
            # Data-center Blackwell (sm_100) must NOT be misclassified as sm_120
            # even when the name contains "Blackwell".
            ("NVIDIA B200", "sm_100"),
            ("NVIDIA B200 (Blackwell)", "sm_100"),
            ("NVIDIA B100", "sm_100"),
        ]
        for name, expected in cases:
            with self.subTest(name=name):
                self.assertEqual(identify_cuda_arch_from_name(name), expected)

    def test_blackwell_consumer(self):
        # Generic "Blackwell" keyword and RTX PRO Blackwell → sm_120
        cases = [
            ("NVIDIA GeForce RTX 5080", "sm_120"),
            ("NVIDIA RTX PRO 3000 Blackwell Generation Laptop GPU", "sm_120"),
            ("NVIDIA RTX PRO 6000 Blackwell", "sm_120"),
            ("NVIDIA RTX PRO 4000 Blackwell", "sm_120"),
        ]
        for name, expected in cases:
            with self.subTest(name=name):
                self.assertEqual(identify_cuda_arch_from_name(name), expected)

    def test_blackwell_datacenter_guard(self):
        # B100/B200 must resolve to sm_100 even when "Blackwell" appears in the name
        cases = [
            ("NVIDIA B100", "sm_100"),
            ("NVIDIA B200", "sm_100"),
            ("NVIDIA B100 Blackwell", "sm_100"),
            ("NVIDIA B200 Blackwell", "sm_100"),
            ("NVIDIA B200 (Blackwell)", "sm_100"),
        ]
        for name, expected in cases:
            with self.subTest(name=name):
                self.assertEqual(identify_cuda_arch_from_name(name), expected)

    def test_non_nvidia_returns_empty(self):
        for name in ["AMD Radeon RX 7900 XTX", "Intel Arc A770", "Apple M3", ""]:
            with self.subTest(name=name):
                self.assertEqual(identify_cuda_arch_from_name(name), "")

    def test_unsupported_nvidia_returns_empty(self):
        # Kepler/Maxwell/Pascal — not in CUDA_SUPPORTED_ARCHS
        for name in ["NVIDIA GTX 1080 Ti", "NVIDIA GTX 980", "NVIDIA Tesla K80"]:
            with self.subTest(name=name):
                result = identify_cuda_arch_from_name(name)
                # May return an sm_XX token but it won't be in the supported set
                if result:
                    self.assertNotIn(result, CUDA_SUPPORTED_ARCHS)


if __name__ == "__main__":
    unittest.main()
