//! NVIDIA discrete GPU energy counter via NVML — `Provenance::HwShunt` while its
//! cumulative-energy counter is live, `Provenance::ModelBased` otherwise.
//!
//! On Volta+ silicon, `nvmlDeviceGetTotalEnergyConsumption()` returns a
//! cumulative-energy counter (mJ since driver reload), summed here across every
//! GPU. That is a vendor energy counter, and its reading is graded a hardware
//! reading — but only while it is demonstrably live for the domain the figure
//! claims (JouleClaw SPEC §1.3): every GPU's counter must keep advancing under
//! the sustained-advance rule ([`crate::CounterLiveness`]) — in at least five
//! sample intervals, four in every five, over at least 300 ms, while the
//! probing thread burns 100 ms of its own CPU time between samples (an RTX
//! 4050 Laptop GPU, driver 595.84, republishes the counter every ~100 ms, so a
//! live counter has moved in nearly every interval; even at a 2 W idle that is
//! ~0.2 J an interval, against a 1 mJ quantum). And every GPU's counter must
//! plausibly count its board's energy over the probe, bounded in **both**
//! directions: no faster than 1.5 times the board's maximum power limit, or
//! than twice the board's own power draw plus 5 W, and no slower than a
//! quarter of that draw once 5 W of slack is taken off it
//! ([`crate::CounterLiveness::Implausible`]). The lower bound is the one that
//! matters at a ceiling — a counter reading low is graded as hardware and then
//! spends a budget it is not measuring — while a counter running fast fails
//! safe. That RTX 4050's counter runs at 4 to 30 times its board's draw while
//! it is polled, so it is declined there. Probed once per process and cached.
//!
//! Where that counter is unsupported (pre-Volta, some consumer parts) or not
//! live, the counter integrates `power_usage()` over wall time instead. That
//! integral is a Riemann sum of the driver's sampled power, not a counter: it
//! advances whenever time passes, so no liveness probe can vouch for it, and
//! its readings are graded [`Provenance::ModelBased`] — the grade JouleClaw
//! SPEC §1.3 gives NVML `power.draw`. Never `HwShunt`.
//!
//! The domain is [`EnergyDomain::Gpu`]: every GPU NVML enumerates, nothing
//! else. The CPU work that drives them is not in it.
//!
//! Resolution: ~1 mJ (the counter's quantum). Minimum useful sampling window:
//! ~100 ms — the counter does not update faster than that on any device
//! tested, so faster polls return the same sampled value (measured: a stale
//! read fraction of 93.4 % at 0.5 ms and 86.3 % at 10 ms, falling to 1.0 % at
//! 100 ms). Accuracy against integrated power above that floor is
//! device-class dependent; see `jouleclaw/SPEC.md` §1.3.
//!
//! Without the `nvml` cargo feature this module is not compiled.

#![forbid(unsafe_code)]

use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::OnceLock;
use std::time::{Duration, SystemTime, UNIX_EPOCH};

use crate::meter::{EnergyMeter, EnergySource};
use crate::{CounterLiveness, EnergyCounter, EnergyDomain, EnergyError, EnergyReading, Provenance};

/// Thread CPU time the liveness probe burns between samples of the
/// cumulative-energy counter.
///
/// Longer than the counter's ~100 ms update period, not equal to it. At a tick
/// equal to the update period roughly half the intervals catch no update, the
/// probe fails [`crate::liveness::LIVE_SHARE`] (four in five) and a perfectly
/// healthy counter is graded not-live. 150 ms matches
/// `sandbox-meter`'s `NVML_PROBE_TICK`, and
/// [`crate::liveness::MIN_LIVE_ADVANCES`] (5) intervals then take 750 ms —
/// inside the 1000 ms [`crate::liveness::LIVENESS_BUDGET`].
const PROBE_TICK: Duration = Duration::from_millis(150);

/// How `read()` obtains a figure, decided once at construction.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Mode {
    /// NVML did not initialise, or found no GPU: no reading.
    Unavailable,
    /// The cumulative-energy counter, live: a hardware reading (`HwShunt`).
    Counter,
    /// The integral of `power_usage()`: a vendor energy model (`ModelBased`).
    PowerIntegral,
}

/// The reading mode for a host whose cumulative-energy counter is `supported`
/// (or not) and whose liveness probe said `liveness`: the counter only when it
/// is supported AND live, the power integral otherwise.
fn mode_for(supported: bool, liveness: CounterLiveness) -> Mode {
    match (supported, liveness) {
        (true, CounterLiveness::Live) => Mode::Counter,
        // An unjudgeable counter is no more admissible than an implausible
        // one: neither can be graded `HwShunt`, so both fall back to the
        // vendor's power model.
        (
            _,
            CounterLiveness::Live
            | CounterLiveness::Frozen
            | CounterLiveness::Unavailable
            | CounterLiveness::Implausible
            | CounterLiveness::Unjudgeable,
        ) => Mode::PowerIntegral,
    }
}

/// The grade of every reading a counter in `mode` produces.
fn grade(mode: Mode) -> Provenance {
    match mode {
        Mode::Counter => Provenance::HwShunt,
        Mode::PowerIntegral => Provenance::ModelBased,
        // No reading is produced; this is the grade a reading would have had
        // from the counter, as the IOReport backend reports for a declined one.
        Mode::Unavailable => Provenance::HwShunt,
    }
}

/// NVIDIA discrete GPU energy counter.
///
/// Initialises NVML at construction time. If NVML init fails or no
/// GPUs are present the counter is permanently unavailable.
pub struct NvmlCounter {
    nvml: Option<nvml_wrapper::Nvml>,
    device_count: u32,
    mode: Mode,
    liveness: CounterLiveness,
    cumulative_uj: AtomicU64,
    last_timestamp_ms: AtomicU64,
}

impl Default for NvmlCounter {
    fn default() -> Self {
        Self::new()
    }
}

impl NvmlCounter {
    /// Initialise NVML and enumerate GPUs. Where every GPU supports the
    /// cumulative-energy counter, the first call in a process runs its liveness
    /// probe on the calling thread (about 0.5 s where it is live, about 1 s where
    /// it is not) and caches the answer ([`Self::liveness`]).
    pub fn new() -> Self {
        let nvml = match nvml_wrapper::Nvml::init() {
            Ok(nvml) => nvml,
            Err(e) => {
                tracing::debug!("NVML init failed: {e}");
                return Self::unavailable(None);
            }
        };
        let device_count = nvml.device_count().unwrap_or(0);
        if device_count == 0 {
            tracing::debug!("NVML initialized but no GPUs found");
            return Self::unavailable(Some(nvml));
        }
        // The cumulative-energy counter must be readable on every GPU the figure
        // sums; `NotSupported` on any of them leaves the power integral.
        let supported = total_energy_uj(&nvml, device_count).is_some();
        let liveness = if supported {
            static LIVENESS: OnceLock<CounterLiveness> = OnceLock::new();
            *LIVENESS.get_or_init(|| probe_liveness(&nvml, device_count))
        } else {
            CounterLiveness::Unavailable
        };
        let mode = mode_for(supported, liveness);
        tracing::info!(gpu_count = device_count, supported, ?liveness, ?mode, "NVML initialized");
        Self::with(Some(nvml), device_count, mode, liveness)
    }

    fn with(nvml: Option<nvml_wrapper::Nvml>, device_count: u32, mode: Mode, liveness: CounterLiveness) -> Self {
        Self {
            nvml,
            device_count,
            mode,
            liveness,
            cumulative_uj: AtomicU64::new(0),
            last_timestamp_ms: AtomicU64::new(0),
        }
    }

    /// An unavailable counter (NVML init failed or no GPUs), keeping any NVML
    /// handle so diagnostics can still query the library.
    fn unavailable(nvml: Option<nvml_wrapper::Nvml>) -> Self {
        Self::with(nvml, 0, Mode::Unavailable, CounterLiveness::Unavailable)
    }

    /// Whether NVML initialised and at least one GPU was found.
    pub fn is_available(&self) -> bool {
        self.mode != Mode::Unavailable
    }

    /// `true` when `read()` uses the hardware's cumulative-energy counter
    /// (`nvmlDeviceGetTotalEnergyConsumption`) — supported on every GPU and
    /// live — and grades its readings `HwShunt`; `false` when it integrates
    /// `power_usage()` and grades them `ModelBased`, or reads nothing.
    pub fn uses_true_counter(&self) -> bool {
        self.mode == Mode::Counter
    }

    /// The cumulative-energy counter's liveness on this host:
    /// [`CounterLiveness::Live`] exactly when [`Self::uses_true_counter`];
    /// `Unavailable` where it is unsupported or NVML has no GPU.
    pub fn liveness(&self) -> CounterLiveness {
        self.liveness
    }
}

/// Every GPU's cumulative energy, in µJ, in device order; `None` if any of them
/// cannot be read.
/// The real liveness probe of every GPU's cumulative-energy counter: the
/// sustained-advance rule at [`PROBE_TICK`], with every GPU's power draw read
/// at each sample and its maximum power limit once, so that a live counter
/// which ran faster than its board can draw, or far slower than it, is
/// [`CounterLiveness::Implausible`] ([`crate::liveness::GpuWatch`]).
fn probe_liveness(nvml: &nvml_wrapper::Nvml, device_count: u32) -> CounterLiveness {
    let ceilings: Vec<Option<u32>> = (0..device_count)
        .map(|i| {
            let limits = nvml.device_by_index(i).ok()?.power_management_limit_constraints().ok()?;
            (limits.max_limit > 0).then_some(limits.max_limit)
        })
        .collect();
    let draws = move || -> Vec<Option<u32>> {
        (0..device_count).map(|i| nvml.device_by_index(i).ok()?.power_usage().ok()).collect()
    };
    let watch = std::cell::RefCell::new(crate::liveness::GpuWatch::default());
    let seen = &watch;
    let verdict = crate::liveness::probe(
        move || {
            total_energy_uj(nvml, device_count).map(|first| {
                (first.len(), move || {
                    let sample = total_energy_uj(nvml, device_count);
                    seen.borrow_mut().saw(sample.as_deref(), &draws(), std::time::Instant::now());
                    sample
                })
            })
        },
        PROBE_TICK,
    );
    watch.into_inner().liveness(verdict, &ceilings)
}

fn total_energy_uj(nvml: &nvml_wrapper::Nvml, device_count: u32) -> Option<Vec<u64>> {
    (0..device_count)
        .map(|i| {
            let mj = nvml.device_by_index(i).ok()?.total_energy_consumption().ok()?;
            Some(mj.saturating_mul(1_000))
        })
        .collect()
}

impl EnergyCounter for NvmlCounter {
    fn domain(&self) -> EnergyDomain {
        EnergyDomain::Gpu
    }

    fn provenance(&self) -> Provenance {
        grade(self.mode)
    }

    fn resolution_uj(&self) -> u64 {
        // ~1 mJ vendor-reported. Conservative.
        1_000
    }

    fn min_window_ns(&self) -> u64 {
        // The counter's update period: ~100 ms on every device tested. Below
        // it a read repeats the preceding value, so a shorter window is not a
        // measurement of that window. This is the noise floor this method
        // names; it is NOT a promise of accuracy against integrated power,
        // which is device-class dependent above it (SPEC §1.3).
        100_000_000
    }

    fn read(&self) -> Result<EnergyReading, EnergyError> {
        let nvml = self
            .nvml
            .as_ref()
            .ok_or(EnergyError::NoCounter(EnergyDomain::Gpu))?;

        if self.mode == Mode::Unavailable {
            return Err(EnergyError::NoCounter(EnergyDomain::Gpu));
        }

        let now_ms = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .map_err(|e| EnergyError::Platform(e.to_string()))?
            .as_millis() as u64;

        // The hardware's own cumulative-energy counter, summed across GPUs, read
        // only where it was observed live. No integration, so no timing error.
        if self.mode == Mode::Counter {
            let mut total_uj: u64 = 0;
            for i in 0..self.device_count {
                let device = nvml
                    .device_by_index(i)
                    .map_err(|e| EnergyError::Platform(format!("GPU {i}: {e}")))?;
                let mj = device
                    .total_energy_consumption()
                    .map_err(|e| EnergyError::Platform(format!("GPU {i} energy: {e}")))?;
                total_uj = total_uj.saturating_add(mj.saturating_mul(1_000));
            }
            return Ok(EnergyReading {
                uj: total_uj,
                timestamp_ns: now_ms.saturating_mul(1_000_000),
                domain: EnergyDomain::Gpu,
                provenance: grade(self.mode),
            });
        }

        // The power integral: the shunt-sampled instantaneous power, integrated
        // over wall time — a Riemann sum, hence ~100 ms windows, graded ModelBased.
        let mut total_milliwatts: u64 = 0;
        for i in 0..self.device_count {
            let device = nvml
                .device_by_index(i)
                .map_err(|e| EnergyError::Platform(format!("GPU {i}: {e}")))?;
            let power_mw = device
                .power_usage()
                .map_err(|e| EnergyError::Platform(format!("GPU {i} power: {e}")))?;
            total_milliwatts = total_milliwatts.saturating_add(power_mw as u64);
        }

        let watts = total_milliwatts as f64 / 1000.0;
        let last_ms = self.last_timestamp_ms.swap(now_ms, Ordering::Relaxed);
        if last_ms > 0 && now_ms > last_ms {
            let elapsed_secs = (now_ms - last_ms) as f64 / 1000.0;
            let delta_uj = (watts * elapsed_secs * 1_000_000.0) as u64;
            self.cumulative_uj.fetch_add(delta_uj, Ordering::Relaxed);
        }

        let total_uj = self.cumulative_uj.load(Ordering::Relaxed);
        Ok(EnergyReading {
            uj: total_uj,
            timestamp_ns: now_ms.saturating_mul(1_000_000),
            domain: EnergyDomain::Gpu,
            provenance: grade(self.mode),
        })
    }
}

impl EnergyMeter for NvmlCounter {
    fn name(&self) -> &'static str {
        "nvml"
    }

    fn is_available(&self) -> bool {
        self.mode != Mode::Unavailable
    }

    fn energy_source(&self) -> EnergySource {
        EnergySource::WallPower
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::liveness::tests::EVERY;
    use crate::liveness::{probe_with, LIVENESS_BUDGET};
    use std::cell::Cell;

    #[test]
    fn nvml_counter_unavailable_without_gpu() {
        // CI / macOS / non-NVIDIA hosts will all hit this branch.
        let c = NvmlCounter::new();
        assert_eq!(c.name(), "nvml");
        assert_eq!(c.domain(), EnergyDomain::Gpu);
        if !c.is_available() {
            assert!(c.read().is_err());
            assert!(!c.uses_true_counter());
            assert_eq!(c.liveness(), CounterLiveness::Unavailable);
        }
        eprintln!("nvml on this host: available={} true_counter={} liveness={:?} grade={:?}",
            c.is_available(), c.uses_true_counter(), c.liveness(), c.provenance());
    }

    /// The liveness probe's tick must CLEAR the counter's update period, not
    /// equal it.
    ///
    /// At a tick equal to the period, roughly half the intervals catch no
    /// update, the probe misses [`crate::liveness::LIVE_SHARE`] (four in five)
    /// and grades a perfectly healthy counter not-live — which demotes a real
    /// `HwShunt` reading to the power integral. The tick was 100 ms against a
    /// period documented as ~50 ms; the period is measured at ~100 ms, so the
    /// margin was zero and the stated rationale ("two of its ~50 ms sample
    /// periods") was false.
    ///
    /// The other half of the constraint: [`crate::liveness::MIN_LIVE_ADVANCES`]
    /// intervals have to fit inside one [`crate::liveness::LIVENESS_BUDGET`],
    /// or the probe runs past its budget on every healthy host.
    #[test]
    fn the_probe_tick_clears_the_update_period_and_fits_the_budget() {
        let c = NvmlCounter::new();
        let tick = u64::try_from(PROBE_TICK.as_nanos()).expect("tick fits u64");
        assert!(
            tick > c.min_window_ns(),
            "probe tick {tick} ns must exceed the counter's {} ns update period",
            c.min_window_ns()
        );
        let needed = PROBE_TICK * crate::liveness::MIN_LIVE_ADVANCES;
        assert!(
            needed <= crate::liveness::LIVENESS_BUDGET,
            "{needed:?} for MIN_LIVE_ADVANCES intervals must fit in {:?}",
            crate::liveness::LIVENESS_BUDGET
        );
    }

    #[test]
    fn nvml_resolution_is_millijoule_floor() {
        let c = NvmlCounter::new();
        // NVML is honest only down to ~1 mJ.
        assert!(c.resolution_uj() >= 1_000);
        // ~100 ms minimum window: the counter's measured update period.
        assert!(c.min_window_ns() >= 100_000_000);
    }

    /// The cumulative-energy counter is read, as a hardware reading, only where
    /// it is supported AND live. A frozen or unprobed counter leaves the power
    /// integral, which is never graded `HwShunt`.
    #[test]
    fn only_a_live_supported_counter_is_a_hardware_reading() {
        use CounterLiveness::*;
        for (supported, liveness, mode) in [
            (true, Live, Mode::Counter),
            (true, Frozen, Mode::PowerIntegral),
            (true, Unavailable, Mode::PowerIntegral),
            (false, Unavailable, Mode::PowerIntegral),
            (false, Live, Mode::PowerIntegral),
            (true, Implausible, Mode::PowerIntegral),
        ] {
            assert_eq!(mode_for(supported, liveness), mode, "supported={supported} {liveness:?}");
            let c = NvmlCounter::with(None, 1, mode, liveness);
            assert_eq!(c.uses_true_counter(), mode == Mode::Counter);
            let want = if mode == Mode::Counter { Provenance::HwShunt } else { Provenance::ModelBased };
            assert_eq!(c.provenance(), want, "supported={supported} {liveness:?}");
            assert!(c.is_available());
        }
        assert_eq!(grade(Mode::PowerIntegral), Provenance::ModelBased, "a power.draw integral is a vendor model");
    }

    /// The probe over scripted per-GPU counters (mJ → µJ, as `total_energy_uj`
    /// returns them) and a clock advancing `EVERY` per tick.
    fn nvml_probe(gpus: usize, mj: impl Fn(usize, usize) -> Option<u64>) -> CounterLiveness {
        let t = Cell::new(Duration::ZERO);
        let n = Cell::new(0usize);
        let sample = || {
            let i = n.get();
            n.set(i + 1);
            (0..gpus).map(|g| mj(i, g).map(|v| v * 1_000)).collect::<Option<Vec<u64>>>()
        };
        probe_with(
            || sample().map(|first| (first.len(), sample)),
            || t.set(t.get() + EVERY),
            || t.get(),
            LIVENESS_BUDGET,
        )
    }

    #[test]
    fn every_gpu_must_keep_advancing() {
        // ~60 W idle per GPU: ~2 J per interval.
        assert_eq!(nvml_probe(2, |i, g| Some(5_000_000 * (g as u64 + 1) + 2_000 * i as u64)), CounterLiveness::Live);
        assert_eq!(
            nvml_probe(2, |i, g| Some(if g == 0 { 5_000_000 + 2_000 * i as u64 } else { 9_000_000 })),
            CounterLiveness::Frozen,
            "one GPU's counter held: the sum would leave its share of the domain unread"
        );
        assert_eq!(nvml_probe(1, |_, _| Some(5_000_000)), CounterLiveness::Frozen);
        assert_eq!(nvml_probe(1, |_, _| None), CounterLiveness::Unavailable);
        assert_eq!(nvml_probe(0, |_, _| Some(1)), CounterLiveness::Frozen, "no GPU is never live");
    }
}