//! Linux RAPL through the energy-status MSRs, against the hardware: on a host
//! whose powercap lists no package zone, the meter's `rapl-msr` probe, read by
//! the meter and by an oracle the tests declare here.
//!
//! In a test binary of its own, so the RAPL probe behind `rapl_status()` runs
//! in this process, and the tests take one lock: each burns CPU and reads
//! machine-wide counters. The oracle reads the MSRs itself — `/proc/cpuinfo`'s
//! vendor and family, each CPU's `topology/physical_package_id` (and `die_id`
//! on Intel), and each domain's power-unit and energy-status MSR through
//! `/dev/cpu/<n>/msr` — and applies the sustained-advance rule on its own.
//! None of these tests passes without asserting: where powercap lists a zone
//! that is or may be a package, the meter must not read the MSRs; where the
//! MSRs cannot be read (no `msr` module, no `CAP_SYS_RAWIO`, a CI container),
//! RAPL must not be live and the declined probe must say why; where they read,
//! the status must agree with the oracle, and a span over a CPU burn must be
//! one `Cpu` record, `HwShunt`, from `rapl-msr`, whose figure sits inside the
//! oracle's bracket. Each test prints one `rapl msr:` line.

#![cfg(target_os = "linux")]

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

use sandbox_meter::probe::{rapl_status, CounterStatus, PowerMeter, RaplPackages};
use sandbox_meter::Provenance;
use sandbox_schema::DeviceClass;

/// One test at a time: each burns CPU and reads machine-wide counters.
static SERIAL: Mutex<()> = Mutex::new(());

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

/// The tests' own reading of the RAPL MSRs, read here and nowhere else, and
/// the sustained-advance verdict over them by the family's rule. Nothing in it
/// comes from the crate under test.
mod oracle {
    use std::collections::BTreeMap;
    use std::os::raw::{c_int, c_long};
    use std::os::unix::fs::FileExt;
    use std::path::{Path, PathBuf};
    use std::time::{Duration, Instant};

    /// Where the kernel lists its powercap zones.
    pub const POWERCAP: &str = "/sys/class/powercap";

    #[repr(C)]
    struct Timespec {
        tv_sec: c_long,
        tv_nsec: c_long,
    }

    unsafe extern "C" {
        fn clock_gettime(clock_id: c_int, tp: *mut Timespec) -> c_int;
    }

    /// `CLOCK_THREAD_CPUTIME_ID` in Linux `<linux/time.h>`.
    const CLOCK_THREAD_CPUTIME_ID: c_int = 3;

    fn thread_cpu() -> Duration {
        let mut ts = Timespec { tv_sec: 0, tv_nsec: 0 };
        assert_eq!(unsafe { clock_gettime(CLOCK_THREAD_CPUTIME_ID, &mut ts) }, 0, "clock_gettime");
        Duration::new(ts.tv_sec as u64, ts.tv_nsec as u32)
    }

    /// Spin this thread until it has used `d` of its own CPU time, or
    /// `wall_cap` of wall time has passed. Returns the CPU it got.
    pub fn burn(d: Duration, wall_cap: Duration) -> Duration {
        let (started, wall) = (thread_cpu(), Instant::now());
        let mut acc: u64 = 0;
        while thread_cpu() - started < d && wall.elapsed() < wall_cap {
            for _ in 0..1_000 {
                acc = std::hint::black_box(acc.wrapping_add(1));
            }
        }
        thread_cpu() - started
    }

    fn digits(s: &str) -> bool {
        !s.is_empty() && s.bytes().all(|b| b.is_ascii_digit())
    }

    /// Whether powercap lists a top-level `intel-rapl:<n>` zone that is a
    /// package (`name` reads `package-<n>`, or `package-<p>-die-<d>` for a die
    /// of a multi-die package) or may be one (`name` does not read): the meter
    /// then reads powercap or nothing, never the MSRs.
    pub fn powercap_may_list_a_package() -> bool {
        let named = |n: &str| match n.trim().strip_prefix("package-") {
            Some(rest) => match rest.split_once("-die-") {
                Some((p, d)) => digits(p) && digits(d),
                None => digits(rest),
            },
            None => false,
        };
        let Ok(entries) = std::fs::read_dir(POWERCAP) else { return false };
        entries.flatten().any(|e| {
            let dir = e.file_name().to_string_lossy().into_owned();
            dir.strip_prefix("intel-rapl:").is_some_and(digits)
                && match std::fs::read_to_string(e.path().join("name")) {
                    Ok(name) => named(&name),
                    Err(_) => true,
                }
        })
    }

    /// One RAPL energy domain as the oracle reads it.
    #[derive(Debug)]
    pub struct Domain {
        /// `package-<p>` (AMD, Hygon) or `package-<p>-die-<d>` (Intel).
        pub name: String,
        pub device: PathBuf,
        pub register: u64,
        /// 2^-`shift` J per count.
        pub shift: u32,
    }

    fn msr(device: &Path, register: u64) -> std::io::Result<u64> {
        let mut raw = [0u8; 8];
        std::fs::File::open(device)?.read_exact_at(&mut raw, register)?;
        Ok(u64::from_le_bytes(raw))
    }

    fn number(p: PathBuf) -> Result<u32, String> {
        let s = std::fs::read_to_string(&p).map_err(|e| format!("{}: {e}", p.display()))?;
        s.trim().parse().map_err(|e| format!("{}: {e}", p.display()))
    }

    /// This host's RAPL domains through its MSRs: Intel family 6 per die of
    /// each package, AMD Zen and Hygon per package, each through the first of
    /// its CPUs whose device reads both MSRs. `Err` when any of that does not
    /// read, or a unit is zero.
    pub fn domains() -> Result<Vec<Domain>, String> {
        let cpuinfo = std::fs::read_to_string("/proc/cpuinfo").map_err(|e| format!("/proc/cpuinfo: {e}"))?;
        let field = |key: &str| {
            cpuinfo.lines().find_map(|l| l.split_once(':').filter(|(k, _)| k.trim() == key).map(|(_, v)| v.trim()))
        };
        let (vendor, family) = (field("vendor_id"), field("cpu family").and_then(|f| f.parse::<u32>().ok()));
        let (unit_reg, energy_reg, per_die) = match (vendor, family) {
            (Some("GenuineIntel"), Some(6)) => (0x606, 0x611, true),
            (Some("AuthenticAMD" | "HygonGenuine"), Some(f)) if f >= 0x17 => (0xC001_0299, 0xC001_029B, false),
            other => return Err(format!("no known RAPL MSRs on {other:?}")),
        };
        let mut by: BTreeMap<(u32, u32), Vec<u32>> = BTreeMap::new();
        for e in std::fs::read_dir("/sys/devices/system/cpu").map_err(|e| e.to_string())?.flatten() {
            let name = e.file_name().to_string_lossy().into_owned();
            let Some(cpu) = name.strip_prefix("cpu").filter(|n| digits(n)).and_then(|n| n.parse::<u32>().ok()) else {
                continue;
            };
            let t = e.path().join("topology");
            let package = number(t.join("physical_package_id"))?;
            let die = if per_die { number(t.join("die_id"))? } else { 0 };
            by.entry((package, die)).or_default().push(cpu);
        }
        if by.is_empty() {
            return Err("no CPU listed".into());
        }
        let mut out = Vec::new();
        for ((package, die), mut cpus) in by {
            cpus.sort_unstable();
            let name = if per_die { format!("package-{package}-die-{die}") } else { format!("package-{package}") };
            let mut first = None;
            let found = cpus.iter().find_map(|cpu| {
                let device = PathBuf::from(format!("/dev/cpu/{cpu}/msr"));
                match msr(&device, unit_reg).and_then(|u| msr(&device, energy_reg).map(|_| u)) {
                    Ok(unit) => Some((device, unit)),
                    Err(e) => {
                        first.get_or_insert(format!("{}: {e}", device.display()));
                        None
                    }
                }
            });
            let Some((device, unit)) = found else {
                return Err(format!("{name}: no CPU's MSRs read (first {})", first.unwrap_or_default()));
            };
            let shift = ((unit >> 8) & 0x1f) as u32;
            if shift == 0 {
                return Err(format!("{name}: a zero energy unit"));
            }
            out.push(Domain { name, device, register: energy_reg, shift });
        }
        Ok(out)
    }

    /// Every domain's raw 32-bit count, in order; `None` if one fails.
    pub fn read(domains: &[Domain]) -> Option<Vec<u32>> {
        domains.iter().map(|d| msr(&d.device, d.register).ok().map(|raw| (raw & 0xffff_ffff) as u32)).collect()
    }

    /// The domains' energy between two reads, in µJ: each count's rise modulo
    /// 2^32 at its own unit, summed, rounded down.
    pub fn energy(domains: &[Domain], before: &[u32], after: &[u32]) -> u64 {
        domains
            .iter()
            .zip(before.iter().zip(after))
            .map(|(d, (&b, &a))| ((u128::from(a.wrapping_sub(b)) * 1_000_000) >> d.shift) as u64)
            .sum()
    }

    /// The sustained-advance run over the domains: sample every count,
    /// burning 25 ms of this thread's CPU (capped at 500 ms of wall time)
    /// between samples, until every domain has advanced in at least five
    /// intervals, at least 80% of those observed, over at least 300 ms; or one
    /// second has passed with at least five intervals; or five seconds have.
    /// Whether it ended live, and each domain's advances out of the intervals.
    pub fn watch(domains: &[Domain]) -> (bool, Vec<u32>, u32) {
        let started = Instant::now();
        let Some(mut prev) = read(domains) else { return (false, Vec::new(), 0) };
        let mut advances = vec![0u32; domains.len()];
        let mut intervals = 0u32;
        let live = |adv: &[u32], n: u32, w: Duration| {
            w >= Duration::from_millis(300) && adv.iter().all(|&a| a >= 5 && u64::from(a) * 5 >= u64::from(n) * 4)
        };
        loop {
            let el = started.elapsed();
            if live(&advances, intervals, el) || el >= Duration::from_secs(5) || (el >= Duration::from_secs(1) && intervals >= 5) {
                break;
            }
            burn(Duration::from_millis(25), Duration::from_millis(500));
            if let Some(now) = read(domains) {
                intervals += 1;
                for ((a, n), p) in advances.iter_mut().zip(&now).zip(&prev) {
                    *a += u32::from(n > p); // a wrap reads as a fall, which is no advance

                }
                prev = now;
            }
        }
        (live(&advances, intervals, started.elapsed()), advances, intervals)
    }
}

/// The RAPL source names in a meter's sources.
fn rapl_sources<'a>(sources: &[&'a str]) -> Vec<&'a str> {
    sources.iter().copied().filter(|s| s.starts_with("rapl-")).collect()
}

/// Where powercap lists a zone that is or may be a package, the meter reads
/// powercap or nothing, never the MSRs. Prints and returns `true` then.
fn powercap_is_read(sources: &[&str]) -> bool {
    if !oracle::powercap_may_list_a_package() {
        return false;
    }
    eprintln!("rapl msr: powercap lists a package zone (or one that may be), so the MSRs are not read; sources {sources:?}");
    assert!(!sources.contains(&"rapl-msr"), "the meter read the MSRs beside a powercap package zone: {sources:?}");
    true
}

/// The status agrees with the oracle's own reading of the MSRs, only a live
/// status opens them and puts `rapl-msr` in the meter, and where they do not
/// read the declined probe says why.
#[test]
fn the_msr_status_agrees_with_the_oracle_and_gates_the_probe() {
    let _one = serial();
    let status = rapl_status();
    let sources = PowerMeter::detect().sources();
    if powercap_is_read(&sources) {
        return;
    }
    let opened = RaplPackages::host();
    match oracle::domains() {
        Err(why) => {
            let declined = opened.as_ref().err().map(|e| e.to_string()).unwrap_or_default();
            eprintln!("rapl msr: the oracle reads no MSR ({why}); status {status:?}; declined: {declined}");
            assert_ne!(status, CounterStatus::Live, "RAPL live although the oracle reads no MSR");
            assert!(rapl_sources(&sources).is_empty(), "{sources:?}");
            assert!(opened.is_err(), "the package counters opened although the oracle reads no MSR");
            if status == CounterStatus::Unavailable {
                assert!(declined.contains("energy-status MSRs"), "the declined probe does not say why: {declined}");
            }
        }
        Ok(domains) => {
            let names: Vec<&str> = domains.iter().map(|d| d.name.as_str()).collect();
            let (live, advances, intervals) = oracle::watch(&domains);
            eprintln!(
                "rapl msr: oracle domains {names:?} live={live} advances {advances:?} of {intervals}; status {status:?}; \
                 sources {sources:?}"
            );
            assert_eq!(status == CounterStatus::Live, live, "the status disagrees with the oracle");
            match opened {
                Ok(p) => {
                    assert!(live, "opened without the oracle seeing every domain live");
                    assert_eq!((p.source(), p.zones()), ("rapl-msr", names), "the probe holds other domains than the oracle's");
                    assert_eq!(rapl_sources(&sources), ["rapl-msr"]);
                }
                Err(e) => {
                    assert!(!live, "live, yet the package counters did not open: {e}");
                    assert!(rapl_sources(&sources).is_empty(), "{sources:?}");
                }
            }
        }
    }
}

/// Over a 300 ms one-thread CPU burn, a span of a meter whose RAPL probe reads
/// the MSRs is one `Cpu` record, `HwShunt`, from `rapl-msr`, its figure inside
/// the oracle's bracket of the same MSRs read around the meter's own reads; and
/// the package counters read through `RaplPackages` sit inside it too. Where
/// the MSRs are not live, no span has a RAPL figure.
#[test]
fn a_cpu_burn_through_the_msrs_reads_inside_the_oracles_bracket() {
    let _one = serial();
    let mut meter = PowerMeter::detect();
    if powercap_is_read(&meter.sources()) {
        return;
    }
    if !meter.sources().contains(&"rapl-msr") {
        eprintln!("rapl msr: no MSR probe in {:?}; status {:?}", meter.sources(), rapl_status());
        assert_ne!(rapl_status(), CounterStatus::Live);
        let ((), m) = meter.measure(|| {
            oracle::burn(Duration::from_millis(50), Duration::from_secs(5));
        });
        assert!(
            !m.devices.iter().any(|d| d.class == DeviceClass::Cpu && d.provenance == Provenance::HwShunt),
            "a CPU record was graded HwShunt with RAPL declined: {:?}",
            m.devices
        );
        return;
    }
    let domains = oracle::domains().expect("the meter reads the MSRs, so the oracle can");
    let slack = domains.len() as u64; // each domain's figure is floored once per read
    let read = || oracle::read(&domains).expect("an oracle read of the MSRs");

    let o0 = read();
    let span = meter.start();
    let o1 = read();
    let cpu = oracle::burn(Duration::from_millis(300), Duration::from_secs(10));
    let o2 = read();
    let m = span.stop(&mut meter);
    let o3 = read();
    let (inner, outer) = (oracle::energy(&domains, &o1, &o2), oracle::energy(&domains, &o0, &o3));
    let records: Vec<_> = m.devices.iter().filter(|d| d.class == DeviceClass::Cpu).collect();
    assert_eq!(records.len(), 1, "not one CPU record: {:?}", m.devices);
    let d = records[0];
    assert_eq!((d.provenance, m.source), (Provenance::HwShunt, "rapl-msr"), "{d:?}");
    let uj = d.measured_uj.expect("an HwShunt record carries its measured figure");

    let mut packages = RaplPackages::host().expect("the meter has the MSR probe, so the package counters open");
    let p0 = read();
    let c0 = packages.cumulative_uj().expect("live MSRs read");
    let p1 = read();
    oracle::burn(Duration::from_millis(300), Duration::from_secs(10));
    let p2 = read();
    let c1 = packages.cumulative_uj().expect("live MSRs read");
    let p3 = read();
    let (p_inner, p_outer) = (oracle::energy(&domains, &p1, &p2), oracle::energy(&domains, &p0, &p3));
    eprintln!(
        "rapl msr: {cpu:?} of thread CPU; span {uj} µJ (oracle inner {inner}, outer {outer}); packages {} µJ (oracle \
         inner {p_inner}, outer {p_outer}); domains {:?}",
        c1 - c0,
        packages.zones()
    );
    assert!(inner > 0, "the MSRs did not advance over the burn");
    assert!(inner.saturating_sub(slack) <= uj && uj <= outer + slack, "the span is outside the oracle's bracket");
    assert!(
        p_inner.saturating_sub(slack) <= c1 - c0 && c1 - c0 <= p_outer + slack,
        "the package counters are outside the oracle's bracket"
    );
    assert!(cpu >= Duration::from_millis(300), "the burn got only {cpu:?} of CPU");
}