// SPDX-License-Identifier: Apache-2.0

//! NVIDIA detection via NVML (`libnvidia-ml`), loaded at runtime from the

//! installed driver. The CUDA toolkit is **not** required and nothing links at

//! build time — `nvml-wrapper` `dlopen`s the library lazily, so a host without

//! an NVIDIA driver simply yields no GPUs.

//!

//! Gated behind the default `nvidia` feature; build with

//! `default-features = false` to drop the `nvml-wrapper` dependency entirely.

#[cfg(feature = "nvidia")]

use crate::{ArchTarget, ComputeCapability, CudaHost, CudaVersion, GpuInfo, Vendor};

#[cfg(feature = "nvidia")]

use nvml_wrapper::Nvml;

#[cfg(feature = "nvidia")]

use std::sync::OnceLock;

/// The process-wide NVML handle.

///

/// NVML is initialized at most once and **never shut down**. This is

/// deliberate, not an oversight: cycling `nvmlInit`/`nvmlShutdown` permanently

/// costs one file descriptor per cycle (an `eventfd` that shutdown does not

/// return), so a caller polling [`detect`](crate::detect) on a timer exhausts

/// its fd table and can no longer `accept()` connections. Holding one handle is

/// flat across calls, and queries against it still return live values — so

/// `memory_info()` readouts stay current.

///

/// Do not add a shutdown or make this handle droppable.

#[cfg(feature = "nvidia")]

static NVML: OnceLock<Nvml> = OnceLock::new();

/// The shared NVML handle, initializing it on first use.

///

/// Only *success* is cached. A failed init allocates no file descriptors, so

/// retrying costs nothing but a failed `dlopen`; caching the failure instead

/// would mean a host whose driver loads after the first call — or a daemon that

/// starts before the driver is up — reports "no NVIDIA GPU" until it restarts.

#[cfg(feature = "nvidia")]

fn nvml() -> Option<&'static Nvml> {

    if let Some(
nvml0
) = NVML.get() {

        return Some(nvml);

    }

    let 
nvml0
 = Nvml::init().ok()?;

    // A concurrent first call may have won the race, in which case our handle

    // is dropped here and theirs is returned. `nvmlInit`/`nvmlShutdown` are

    // reference counted, so the winner's handle stays valid. This costs one fd,

    // once, and only when two threads make the very first call simultaneously.

    Some(NVML.get_or_init(|| nvml))

}

/// Read a device's compute capability, rejecting values NVML shouldn't produce.

///

/// NVML reports these as signed and a misbehaving driver can return negatives.

/// Reject them here rather than casting them into huge unsigned values and

/// making every caller re-derive the check.

#[cfg(feature = "nvidia")]

fn device_compute_capability(device: &nvml_wrapper::Device<'_>) -> Option<ComputeCapability> {

    let capability = device.cuda_compute_capability().ok()?;

    Some(ComputeCapability {

        major: u32::try_from(capability.major).ok()?,

        minor: u32::try_from(capability.minor).ok()?,

    })

}

#[cfg(feature = "nvidia")]

pub(crate) fn detect() -> Vec<GpuInfo> {

    let Some(
nvml0
) = nvml() else {

        return Vec::new();

    };

    let Ok(count) = nvml.device_count() else {

        return Vec::new();

    };

    let mut gpus = Vec::new();

    for index in 0..count {

        let Ok(device) = nvml.device_by_index(index) else {

            continue;

        };

        let Ok(memory) = device.memory_info() else {

            continue;

        };

        gpus.push(GpuInfo {

            name: device.name().unwrap_or_else(|_| "NVIDIA GPU".to_string()),

            vendor: Vendor::Nvidia,

            total_bytes: memory.total,

            free_bytes: Some(memory.free),

            used_bytes: Some(memory.used),

            arch_target: device_compute_capability(&device).map(ArchTarget::Sm),

        });

    }

    gpus

}

/// Host-wide CUDA properties, read from the same shared NVML handle as

/// [`detect`] so the process keeps exactly one NVML owner.

#[cfg(feature = "nvidia")]

pub(crate) fn cuda_host() -> Option<CudaHost> {

    let 
nvml0
 = nvml()?;

    // Compute capability is per-device — and is also reported that way on

    // `GpuInfo` — but consumers use this struct to choose one build target for

    // the host, so device 0 is the meaningful answer here.

    let compute_capability = device_compute_capability(&nvml.device_by_index(0).ok()?)?;

    let packed = nvml.sys_cuda_driver_version().ok()?;

    if packed <= 0 {

        return None;

    }

    // NVML packs this as `major * 1000 + minor * 10`; let the binding unpack it

    // rather than open-coding the arithmetic.

    let driver_version = CudaVersion {

        major: u32::try_from(nvml_wrapper::cuda_driver_version_major(packed)).ok()?,

        minor: u32::try_from(nvml_wrapper::cuda_driver_version_minor(packed)).ok()?,

    };

    Some(CudaHost {

        compute_capability,

        driver_version,

    })

}

#[cfg(not(feature = "nvidia"))]

pub(crate) fn detect() -> Vec<crate::GpuInfo> {

    Vec::new()

}

#[cfg(not(feature = "nvidia"))]

pub(crate) fn cuda_host() -> Option<crate::CudaHost> {

    None

}