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

import sys
import warnings

import pytest

from cuda.bindings import nvml
from cuda_python_test_helpers import driver_version_less_than


def assert_nvml_is_initialized():
    assert nvml.device_get_count_v2() > 0


def assert_nvml_is_uninitialized():
    with pytest.raises(nvml.UninitializedError):
        nvml.device_get_count_v2()


def test_devices_are_the_same_architecture(all_devices):
    # The tests in this directory that use `unsupported_before` will generally
    # skip the entire test after the first device that isn't supported is found.
    # This means that if subsequent devices are of a different architecture,
    # they won't be tested properly.  This tests for the (hopefully rare) case
    # where a system has devices of different architectures and produces a warning.

    def get_architecture_name(arch):
        try:
            return nvml.DeviceArch(arch).name
        except ValueError:
            return f"UNKNOWN_ARCH_ID({arch})"

    all_arches = {nvml.device_get_architecture(device) for device in all_devices}

    if len(all_arches) > 1:
        warnings.warn(
            f"System has devices of multiple architectures ({', '.join(get_architecture_name(x) for x in all_arches)}). "
            f" Some tests may be skipped unexpectedly",
            UserWarning,
        )


@pytest.mark.skipif(sys.platform == "win32", reason="Test not supported on Windows")
@pytest.mark.thread_unsafe(reason="nvml init affects other threads")
@pytest.mark.skipif(not driver_version_less_than(13040), reason="Init behavior changed in CUDA 13.4")
def test_init_ref_count():
    """
    Verifies that we can call NVML shutdown and init(2) multiple times, and that ref counting works
    """
    with pytest.raises(nvml.UninitializedError):
        nvml.shutdown()

    assert_nvml_is_uninitialized()

    for i in range(3):
        # Init 5 times
        for j in range(5):
            nvml.init_v2()
            assert_nvml_is_initialized()

        # Shutdown 4 times, NVML should remain initailized
        for j in range(4):
            nvml.shutdown()
            assert_nvml_is_initialized()

        # Shutdown the final time
        nvml.shutdown()
        assert_nvml_is_uninitialized()


def test_init_check_index(nvml_init):
    """
    Verifies that the index from nvmlDeviceGetIndex is correct
    """
    dev_count = nvml.device_get_count_v2()
    for idx in range(dev_count):
        handle = nvml.device_get_handle_by_index_v2(idx)
        # Verify that the index matches
        assert idx == nvml.device_get_index(handle)
