// 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::{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))

}

#[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),

        });

    }

    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, but consumers use it to choose a build

    // target for the host, so device 0 is the meaningful answer.

    let capability = nvml

        .device_by_index(0)

        .ok()?

        .cuda_compute_capability()

        .ok()?;

    // 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.

    let compute_capability = ComputeCapability {

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

        minor: u32::try_from(capability.minor).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

}