//! NVML against the hardware (feature `nvml`): the GPUs' total-energy
//! counters, read by the meter and by an oracle the tests declare here.
//!
//! In a test binary of its own, so the NVML probe runs in this process, and
//! the tests take one lock. The oracle opens `libnvidia-ml` itself, measures
//! how often the counter is republished, samples it at the probe's own tick
//! (100 ms) by the sustained-advance rule and holds each counter to its
//! board's power ceiling, which it asks of `nvidia-smi`. None of these tests
//! passes without asserting: where there is no NVML (a CI container has no
//! GPU), the counter must say so and agree with the oracle; where the
//! counter is live, both must say so, and its readings must be `HwShunt` over
//! the `Gpu` domain — the GPU, never a whole-host or CPU figure — no faster
//! than the board can draw; where it is not live, its readings are the
//! power-draw integral, `ModelBased`, never `HwShunt`.

#![cfg(all(target_os = "linux", feature = "nvml"))]

use std::sync::{Mutex, MutexGuard};
use std::time::Duration;

use jouleclaw_energy::nvml::NvmlCounter;
use jouleclaw_energy::{CounterLiveness, EnergyCounter, EnergyDomain, Provenance};

/// One test at a time: each reads machine-wide counters, and the oracle's
/// own sampling must not overlap the meter's.
static SERIAL: Mutex<()> = Mutex::new(());

fn serial() -> MutexGuard<'static, ()> {
    SERIAL.lock().unwrap_or_else(|poisoned| poisoned.into_inner())
}

/// The NVML probe's tick (`PROBE_TICK` in `nvml.rs`): 100 ms of the probing
/// thread's CPU between samples, so at least 100 ms of wall time.
const TICK: Duration = Duration::from_millis(100);

/// The tests' own reading of NVIDIA's total-energy counter: `libnvidia-ml`
/// opened here with `dlopen` and called through symbols looked up here, not
/// through the crate under test or its bindings, and each GPU's power limit
/// and draw asked of `nvidia-smi`, a separate process. It measures how often
/// the counter is republished, and judges it by the family's rule: every
/// GPU's counter advanced in at least five intervals, at least 80% of them,
/// over at least 300 ms, sampled at the crate's own tick; and it is plausibly
/// counting its board's energy: over the window it ran no faster than 1.5
/// times the board's maximum power limit, the board's own power draw (as
/// `nvidia-smi` reports it just before and just after) was itself within that
/// bound, and the counter ran no faster than twice that draw plus 5 W (the
/// draw's documented accuracy). A counter that runs several times faster than
/// its board draws is not counting that board's energy.
#[allow(dead_code)]
mod nvml_oracle {
    use std::ffi::{c_char, c_int, c_uint, c_void, CString};
    use std::time::{Duration, Instant};

    unsafe extern "C" {
        fn dlopen(filename: *const c_char, flag: c_int) -> *mut c_void;
        fn dlsym(handle: *mut c_void, symbol: *const c_char) -> *mut c_void;
    }

    /// `RTLD_NOW` in `<dlfcn.h>`.
    const RTLD_NOW: c_int = 2;

    type Device = *mut c_void;
    type EnergyFn = unsafe extern "C" fn(Device, *mut u64) -> c_int;
    type LimitsFn = unsafe extern "C" fn(Device, *mut c_uint, *mut c_uint) -> c_int;
    type MilliwattsFn = unsafe extern "C" fn(Device, *mut c_uint) -> c_int;

    /// NVML, opened by the oracle.
    pub struct Nvml {
        energy: EnergyFn,
        limits: Option<LimitsFn>,
        power: Option<MilliwattsFn>,
        /// Every GPU whose total-energy counter reads, in index order.
        pub devices: Vec<Device>,
    }

    fn sym(lib: *mut c_void, name: &str) -> Option<*mut c_void> {
        let name = CString::new(name).ok()?;
        let p = unsafe { dlsym(lib, name.as_ptr()) };
        (!p.is_null()).then_some(p)
    }

    /// `libnvidia-ml.so.1`, initialised, with every GPU whose counter reads;
    /// `None` where there is no library, no GPU, or no GPU with the counter.
    pub fn open() -> Option<Nvml> {
        let file = CString::new("libnvidia-ml.so.1").ok()?;
        let lib = unsafe { dlopen(file.as_ptr(), RTLD_NOW) };
        if lib.is_null() {
            return None;
        }
        unsafe {
            let init: unsafe extern "C" fn() -> c_int = std::mem::transmute(sym(lib, "nvmlInit_v2")?);
            let count: unsafe extern "C" fn(*mut c_uint) -> c_int = std::mem::transmute(sym(lib, "nvmlDeviceGetCount_v2")?);
            let by_index: unsafe extern "C" fn(c_uint, *mut Device) -> c_int =
                std::mem::transmute(sym(lib, "nvmlDeviceGetHandleByIndex_v2")?);
            let energy: EnergyFn = std::mem::transmute(sym(lib, "nvmlDeviceGetTotalEnergyConsumption")?);
            let limits = sym(lib, "nvmlDeviceGetPowerManagementLimitConstraints").map(|p| std::mem::transmute::<_, LimitsFn>(p));
            let power = sym(lib, "nvmlDeviceGetPowerUsage").map(|p| std::mem::transmute::<_, MilliwattsFn>(p));
            if init() != 0 {
                return None;
            }
            let mut n: c_uint = 0;
            if count(&mut n) != 0 {
                return None;
            }
            let mut devices = Vec::new();
            for i in 0..n {
                let mut d: Device = std::ptr::null_mut();
                let mut mj = 0u64;
                if by_index(i, &mut d) == 0 && !d.is_null() && energy(d, &mut mj) == 0 {
                    devices.push(d);
                }
            }
            (!devices.is_empty()).then_some(Nvml { energy, limits, power, devices })
        }
    }

    impl Nvml {
        /// Every GPU's total-energy counter, in mJ; `None` if one fails.
        pub fn energy_mj(&self) -> Option<Vec<u64>> {
            self.devices
                .iter()
                .map(|&d| {
                    let mut mj = 0u64;
                    (unsafe { (self.energy)(d, &mut mj) } == 0).then_some(mj)
                })
                .collect()
        }

        /// GPU `i`'s maximum power-management limit, in mW, as NVML gives it.
        pub fn max_limit_mw(&self, i: usize) -> Option<u32> {
            let (mut lo, mut hi) = (0, 0);
            (unsafe { (self.limits?)(self.devices[i], &mut lo, &mut hi) } == 0 && hi > 0).then_some(hi)
        }

        /// GPU `i`'s power draw, in mW, as NVML gives it.
        pub fn power_mw(&self, i: usize) -> Option<u32> {
            let mut mw = 0;
            (unsafe { (self.power?)(self.devices[i], &mut mw) } == 0).then_some(mw)
        }
    }

    /// One GPU as `nvidia-smi` gives it, in watts: (maximum power limit,
    /// enforced limit, draw), each `None` where it gave none.
    pub type SmiPower = (Option<f64>, Option<f64>, Option<f64>);

    /// `nvidia-smi`'s view of every GPU: (maximum power limit, enforced limit,
    /// draw), in watts; `None` without `nvidia-smi`.
    pub fn smi() -> Option<Vec<SmiPower>> {
        let out = std::process::Command::new("nvidia-smi")
            .args(["--query-gpu=power.max_limit,enforced.power.limit,power.draw", "--format=csv,noheader,nounits"])
            .output()
            .ok()?;
        out.status.success().then(|| {
            String::from_utf8_lossy(&out.stdout)
                .lines()
                .map(|l| {
                    let f: Vec<Option<f64>> = l.split(',').map(|v| v.trim().parse().ok()).collect();
                    (f.first().copied().flatten(), f.get(1).copied().flatten(), f.get(2).copied().flatten())
                })
                .collect()
        })
    }

    /// GPU `i`'s ceiling in watts: `nvidia-smi`'s maximum power limit, else
    /// NVML's.
    pub fn ceiling_w(nvml: &Nvml, smi: &Option<Vec<SmiPower>>, i: usize) -> Option<f64> {
        smi.as_ref()
            .and_then(|s| s.get(i))
            .and_then(|(max, _, _)| *max)
            .or_else(|| nvml.max_limit_mw(i).map(|mw| f64::from(mw) / 1000.0))
    }

    /// How often GPU 0's counter is republished: the median gap between its
    /// changes, polled every 2 ms for up to 1.5 s (`None` if it changed
    /// fewer than 6 times).
    pub fn update_period(nvml: &Nvml) -> Option<Duration> {
        let started = Instant::now();
        let mut last = nvml.energy_mj()?[0];
        let mut changes = Vec::new();
        while started.elapsed() < Duration::from_millis(1500) && changes.len() < 16 {
            std::thread::sleep(Duration::from_millis(2));
            let now = nvml.energy_mj()?[0];
            if now != last {
                changes.push(started.elapsed());
                last = now;
            }
        }
        let mut gaps: Vec<Duration> = changes.windows(2).map(|w| w[1] - w[0]).collect();
        if gaps.len() < 5 {
            return None;
        }
        gaps.sort();
        Some(gaps[gaps.len() / 2])
    }

    /// The oracle's verdict on the counters.
    #[derive(Debug, Clone, Copy, PartialEq, Eq)]
    pub enum Verdict {
        /// No NVML, no GPU, or no GPU whose total-energy counter reads.
        Unavailable,
        /// A counter did not keep advancing.
        Frozen,
        /// Every counter kept advancing, but one ran faster than its board
        /// can draw (see the module doc).
        Implausible,
        /// Every counter kept advancing, and one could not be judged because
        /// its board's own telemetry read past the board's own limit. The
        /// yardstick is broken, not necessarily the counter.
        Unjudgeable,
        /// Every counter kept advancing, plausibly its board's energy.
        Live,
    }

    /// The family's plausibility rule for one GPU over a window: its counter's
    /// average `rate_w`, the board's draws `draws_w` around the window and
    /// its maximum power limit `ceiling_w`. A bound whose input is missing
    /// is not applied.
    ///
    /// A draw reading past the board's own limit is not a reading — an RTX
    /// 4050 Laptop GPU reports about 590 W for a second after it wakes — so it
    /// is dropped, and only it. Only when nothing usable is left is there no
    /// yardstick at all, and then the counter is `Unjudgeable` rather than
    /// `Implausible`: the yardstick is broken, not necessarily the counter.
    ///
    /// This oracle is written from the rule rather than from the crate, in
    /// f64 and from `nvidia-smi` rather than the crate's own sampling, so it
    /// is an independent check on the same rule — which is only true while it
    /// encodes the rule the crate implements.
    pub fn judge(rate_w: f64, draws_w: &[f64], ceiling_w: Option<f64>) -> Verdict {
        let cap = ceiling_w.map(|c| 1.5 * c);
        let (usable, unusable): (Vec<f64>, Vec<f64>) =
            draws_w.iter().copied().partition(|d| cap.is_none_or(|c| *d <= c));
        if let Some(c) = cap {
            if rate_w > c {
                return Verdict::Implausible;
            }
            if usable.is_empty() && !unusable.is_empty() {
                return Verdict::Unjudgeable;
            }
        }
        if !usable.is_empty() {
            let mean = usable.iter().sum::<f64>() / usable.len() as f64;
            if rate_w > 2.0 * mean + 5.0 {
                return Verdict::Implausible;
            }
        }
        Verdict::Live
    }

    /// What the oracle saw.
    #[derive(Debug)]
    pub struct Watch {
        pub verdict: Verdict,
        pub intervals: u32,
        pub advances: Vec<u32>,
        pub window: Duration,
        /// Each GPU's average rate over the window, in watts.
        pub rate_w: Vec<f64>,
        /// Each GPU's draw just before and just after the window, in watts.
        pub draws_w: Vec<Vec<f64>>,
        /// Each GPU's ceiling, in watts.
        pub ceiling_w: Vec<Option<f64>>,
    }

    impl Watch {
        /// The verdict, each GPU's advances, its counter's rate, the board's
        /// draw and its maximum limit, for the test's log line.
        pub fn summary(&self) -> String {
            format!(
                "{:?} (advanced {:?} of {} intervals over {} ms; counter rate {:?} W, board draw {:?} W, max power limit {:?} W)",
                self.verdict,
                self.advances,
                self.intervals,
                self.window.as_millis(),
                self.rate_w.iter().map(|w| (w * 10.0).round() / 10.0).collect::<Vec<_>>(),
                self.draws_w,
                self.ceiling_w
            )
        }
    }

    /// The oracle's run over the first `gpus` GPUs (all, if more), sleeping
    /// `tick` between samples, by the sustained-advance rule (stopping when
    /// live, or after a second with at least five intervals, or five
    /// seconds); then each counter's plausibility over the whole window.
    pub fn watch(tick: Duration, gpus: usize) -> Watch {
        let none = Watch {
            verdict: Verdict::Unavailable,
            intervals: 0,
            advances: Vec::new(),
            window: Duration::ZERO,
            rate_w: Vec::new(),
            draws_w: Vec::new(),
            ceiling_w: Vec::new(),
        };
        let Some(nvml) = open() else { return none };
        let n = nvml.devices.len().min(gpus);
        let before = smi();
        let read = || nvml.energy_mj().map(|v| v[..n].to_vec());
        let started = Instant::now();
        let Some(first) = read() else { return none };
        let (mut prev, mut last_at) = (first.clone(), Duration::ZERO);
        let mut advances = vec![0u32; n];
        let mut intervals = 0u32;
        let live = |adv: &[u32], k: u32, w: Duration| {
            w >= Duration::from_millis(300) && adv.iter().all(|&a| a >= 5 && u64::from(a) * 5 >= u64::from(k) * 4)
        };
        loop {
            let el = started.elapsed();
            if live(&advances, intervals, last_at) || el >= Duration::from_secs(5) || (el >= Duration::from_secs(1) && intervals >= 5) {
                break;
            }
            std::thread::sleep(tick);
            if let Some(now) = read() {
                last_at = started.elapsed();
                intervals += 1;
                for ((a, x), p) in advances.iter_mut().zip(&now).zip(&prev) {
                    *a += u32::from(x > p);
                }
                prev = now;
            }
        }
        let after = smi();
        let window = last_at;
        let rate_w: Vec<f64> =
            first.iter().zip(&prev).map(|(a, b)| b.saturating_sub(*a) as f64 / 1000.0 / window.as_secs_f64().max(1e-9)).collect();
        let draw = |s: &Option<Vec<SmiPower>>, i: usize| s.as_ref().and_then(|s| s.get(i)).and_then(|t| t.2);
        let draws_w: Vec<Vec<f64>> = (0..n).map(|i| [draw(&before, i), draw(&after, i)].into_iter().flatten().collect()).collect();
        let ceiling_w: Vec<Option<f64>> = (0..n).map(|i| ceiling_w(&nvml, &after, i)).collect();
        // The verdict that speaks for the run is the most severe of its GPUs',
        // as the crate's is: a counter proven to over-run is a finding, while
        // one board's broken telemetry only withholds a judgement.
        let judged: Vec<Verdict> = (0..n).map(|i| judge(rate_w[i], &draws_w[i], ceiling_w[i])).collect();
        let verdict = if !live(&advances, intervals, window) {
            Verdict::Frozen
        } else if judged.contains(&Verdict::Implausible) {
            Verdict::Implausible
        } else if judged.contains(&Verdict::Unjudgeable) {
            Verdict::Unjudgeable
        } else {
            Verdict::Live
        };
        Watch { verdict, intervals, advances, window, rate_w, draws_w, ceiling_w }
    }

    /// One line on the GPUs for the log: count, counters, ceilings, draw, and
    /// the counter's measured republication `period`.
    pub fn listing(period: Option<Duration>) -> String {
        let Some(nvml) = open() else { return "no NVML device with a total-energy counter".to_string() };
        format!(
            "{} GPU(s) with a total-energy counter, counters {:?} mJ; NVML max power limit {:?} mW, power draw {:?} mW; \
             nvidia-smi (max limit, enforced limit, draw) W {:?}; counter republished every {:?}",
            nvml.devices.len(),
            nvml.energy_mj(),
            (0..nvml.devices.len()).map(|i| nvml.max_limit_mw(i)).collect::<Vec<_>>(),
            (0..nvml.devices.len()).map(|i| nvml.power_mw(i)).collect::<Vec<_>>(),
            smi(),
            period
        )
    }

    /// GPU 0's counter's republication period, where there is one.
    pub fn period() -> Option<Duration> {
        open().and_then(|n| update_period(&n))
    }
}

use nvml_oracle::Verdict;

/// The liveness agrees with the oracle, judged at the probe's own tick, and
/// the tick is no faster than the counter is republished; the grade follows
/// the liveness.
#[test]
fn the_nvml_liveness_agrees_with_the_oracle_at_its_own_tick() {
    let _one = serial();
    let period = nvml_oracle::period();
    let counter = NvmlCounter::new();
    let watch = nvml_oracle::watch(TICK, usize::MAX);
    let liveness = counter.liveness();
    eprintln!(
        "nvml hardware: {}; oracle at a {TICK:?} tick {}; liveness={liveness:?} available={} true_counter={} grade={:?}",
        nvml_oracle::listing(period),
        watch.summary(),
        counter.is_available(),
        counter.uses_true_counter(),
        counter.provenance()
    );
    if let Some(p) = period {
        assert!(TICK >= p * 9 / 10, "the {TICK:?} tick is faster than the counter's {p:?} republication");
    }
    assert_eq!(liveness == CounterLiveness::Live, watch.verdict == Verdict::Live, "the liveness disagrees with the oracle");
    match watch.verdict {
        Verdict::Unavailable => assert_eq!(liveness, CounterLiveness::Unavailable),
        Verdict::Frozen => assert_eq!(liveness, CounterLiveness::Frozen),
        Verdict::Implausible => assert_eq!(liveness, CounterLiveness::Implausible),
        Verdict::Unjudgeable => assert_eq!(liveness, CounterLiveness::Unjudgeable),
        Verdict::Live => {}
    }
    assert_eq!(counter.uses_true_counter(), liveness == CounterLiveness::Live, "the counter mode disagrees");
    assert_eq!(counter.domain(), EnergyDomain::Gpu);
    if counter.is_available() {
        let want = if liveness == CounterLiveness::Live { Provenance::HwShunt } else { Provenance::ModelBased };
        assert_eq!(counter.provenance(), want, "the grade does not follow the liveness");
    }
}

/// Over a one-second span the counter's readings are graded by its liveness
/// over the `Gpu` domain: a live counter's are `HwShunt`, advanced, inside
/// the oracle's own reads around the span, and no faster than the boards can
/// draw; any other's are `ModelBased` (the power-draw integral), never
/// `HwShunt`.
#[test]
fn a_gpu_counter_reads_the_gpu_alone_graded_by_its_liveness() {
    let _one = serial();
    let counter = NvmlCounter::new();
    if !counter.is_available() {
        eprintln!("nvml span: no NVML ({:?})", counter.liveness());
        assert!(counter.read().is_err());
        assert!(nvml_oracle::open().is_none(), "the oracle opens NVML, yet the counter is unavailable");
        return;
    }
    let oracle = nvml_oracle::open();
    let o0 = oracle.as_ref().and_then(|o| o.energy_mj());
    let r0 = counter.read().expect("an available counter reads");
    std::thread::sleep(Duration::from_secs(1));
    let r1 = counter.read().expect("an available counter reads");
    let o1 = oracle.as_ref().and_then(|o| o.energy_mj());
    eprintln!("nvml span: liveness {:?}; readings {r0:?} -> {r1:?}; oracle {o0:?} -> {o1:?} mJ", counter.liveness());
    for r in [&r0, &r1] {
        assert_eq!(r.domain, EnergyDomain::Gpu, "{r:?}");
    }
    if counter.liveness() != CounterLiveness::Live {
        assert!([&r0, &r1].iter().all(|r| r.provenance == Provenance::ModelBased), "a counter that is not live read HwShunt");
        return;
    }
    assert!([&r0, &r1].iter().all(|r| r.provenance == Provenance::HwShunt));
    let oracle = oracle.expect("a live counter, yet the oracle cannot open NVML");
    let (o0, o1) = (o0.expect("oracle read"), o1.expect("oracle read"));
    let bracket: u64 = o0.iter().zip(&o1).map(|(a, b)| b - a).sum::<u64>() * 1_000;
    let uj = r1.uj - r0.uj;
    assert!(uj > 0 && uj <= bracket, "{uj} µJ outside the oracle's (0, {bracket}] µJ");
    let smi = nvml_oracle::smi();
    let ceiling: f64 = (0..oracle.devices.len()).map(|i| nvml_oracle::ceiling_w(&oracle, &smi, i).unwrap_or(f64::INFINITY)).sum();
    let watts = uj as f64 / 1e6 / ((r1.timestamp_ns - r0.timestamp_ns) as f64 / 1e9);
    assert!(watts <= 1.5 * ceiling, "the GPU readings run at {watts:.1} W, past the boards' {ceiling} W ceiling");
}