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


_NVLINK_VERSION_MAPPING = {
    nvml.NvlinkVersion.VERSION_1_0: (1, 0),
    nvml.NvlinkVersion.VERSION_2_0: (2, 0),
    nvml.NvlinkVersion.VERSION_2_2: (2, 2),
    nvml.NvlinkVersion.VERSION_3_0: (3, 0),
    nvml.NvlinkVersion.VERSION_3_1: (3, 1),
    nvml.NvlinkVersion.VERSION_4_0: (4, 0),
    nvml.NvlinkVersion.VERSION_5_0: (5, 0),
}

_NVLINK_VERSION_6_0 = getattr(nvml.NvlinkVersion, "VERSION_6_0", None)
if _NVLINK_VERSION_6_0 is not None:
    _NVLINK_VERSION_MAPPING[_NVLINK_VERSION_6_0] = (6, 0)


class _NvlinkInfoMeta(type):
    @property
    def max_links(cls):
        """
        The statically-defined maximum number of Nvlinks available.  Defined in
        upstream NVML as ``NVML_NVLINK_MAX_LINKS``.

        To find the actual number of Nvlinks available on a device, use
        :py:attr:`Device.get_nvlink_count`.

        .. version-deprecated:: 1.1.0
            This property is deprecated and will be removed in a future release.
            Use :py:attr:`Device.get_nvlink_count` instead.
        """
        warnings.warn(
            "The `max_links` property is deprecated and will be removed in a future release. "
            "Use `Device.get_nvlink_count` instead.",
            DeprecationWarning,
        )

        # This will always return 18, even on CTK 13.3 where it should be 36.
        return nvml.NVLINK_MAX_LINKS


cdef class _NvlinkInfo:
    """
    Nvlink information for a device.
    """
    cdef Device _device
    cdef int _link

    def __init__(self, device: Device, link: int):
        self._device = device
        self._link = link

    @property
    def version(self) -> tuple[int, int]:
        """
        Retrieves the NvLink version for the device and link.

        For all products with NvLink support.

        Returns
        -------
        tuple[int, int]
            The Nvlink version as a tuple of (major, minor).
        """
        version = nvml.device_get_nvlink_version(self._device._handle, self._link)
        if version == nvml.NvlinkVersion.VERSION_INVALID:
            raise RuntimeError("Invalid NvLink version returned for device")
        try:
            return _NVLINK_VERSION_MAPPING[version]
        except KeyError:
            raise RuntimeError(f"Unknown NvLink version {version} returned for device") from None

    @property
    def state(self) -> bool:
        """
        Retrieves the state of the device's Nvlink for the device and link specified.

        For Pascal™ or newer fully supported devices.

        For all products with Nvlink support.

        Returns
        -------
        bool
            `True` if the Nvlink is active.
        """
        return (
            nvml.device_get_nvlink_state(self._device._handle, self._link) == nvml.EnableState.FEATURE_ENABLED
        )


class NvlinkInfo(_NvlinkInfo, metaclass=_NvlinkInfoMeta):
    pass
