//! The energy of an encode ([platform-encode] §8): a counter read around each of an odd
//! number of rounds of the same encode, the median round sealed with its acquisition
//! class and a declared uncertainty that covers the rounds' spread. No counter is
//! divided among its users by busy time.
//!
//! | counter | reads | attribution | fabrics |
//! |---|---|---|---|
//! | `task` | the operating system's account of this process's energy, with the work other processes did for it billed to it (macOS) | `task` | `cpu` |
//! | `rapl` | the processor package zones' energy counters, summed (Linux) | `bracket` | `cpu` |
//! | `gpu-board` | the discrete GPU driver's total-energy counter for one device (Linux) | `bracket` | `gpu`, `media` |
//! | `rapl-and-gpu-board` | both, summed | `bracket` | `cpu`, `gpu`, `media` |
//!
//! A `bracket` figure is energy-measurement §3's marginal: the busy window's energy less
//! the mean power of an idle window, read just before it, over the busy window's length.
//! Every figure is `Partial`: no counter here sees all of a host's energy, and the
//! `task` counter sees none of a hardware engine's.
//!
//! [platform-encode]: ../../../../extensions/platform-encode.md

use super::record::Energy;
use std::time::{Duration, Instant};

/// A counter that only ever increases, in nanojoules.
pub trait Counter {
    fn name(&self) -> &'static str;
    fn attribution(&self) -> &'static str;
    fn fabrics(&self) -> Vec<&'static str>;
    fn read_nj(&mut self) -> Option<u64>;
}

/// Open the counter named `name`.
pub fn open(name: &str) -> Result<Box<dyn Counter>, String> {
    match name {
        #[cfg(target_vendor = "apple")]
        "task" => Ok(Box::new(task::Task)),
        #[cfg(target_os = "linux")]
        "rapl" => linux::Rapl::open().map(|c| Box::new(c) as Box<dyn Counter>),
        #[cfg(target_os = "linux")]
        "gpu-board" => linux::Board::open().map(|c| Box::new(c) as Box<dyn Counter>),
        #[cfg(target_os = "linux")]
        "rapl-and-gpu-board" => Ok(Box::new(linux::Both(linux::Rapl::open()?, linux::Board::open()?))),
        n => Err(format!("no counter {n} on this platform")),
    }
}

#[cfg(target_vendor = "apple")]
mod task {
    use std::ffi::{c_int, c_void};

    unsafe extern "C" {
        fn proc_pid_rusage(pid: c_int, flavor: c_int, buffer: *mut c_void) -> c_int;
        fn getpid() -> c_int;
    }
    const RUSAGE_INFO_V6: c_int = 6;
    // Byte offsets in `struct rusage_info_v6`: ri_billed_energy, ri_serviced_energy,
    // ri_energy_nj.
    const BILLED: usize = 264;
    const SERVICED: usize = 272;
    const ENERGY: usize = 336;

    pub struct Task;

    impl super::Counter for Task {
        fn name(&self) -> &'static str {
            "task"
        }
        fn attribution(&self) -> &'static str {
            "task"
        }
        fn fabrics(&self) -> Vec<&'static str> {
            vec!["cpu"]
        }
        fn read_nj(&mut self) -> Option<u64> {
            let mut buf = [0u64; 128];
            // SAFETY: the buffer is larger than the V6 structure.
            let r = unsafe { proc_pid_rusage(getpid(), RUSAGE_INFO_V6, buf.as_mut_ptr().cast()) };
            if r != 0 {
                return None;
            }
            let at = |off: usize| buf[off / 8];
            Some((at(ENERGY) + at(BILLED)).saturating_sub(at(SERVICED)))
        }
    }
}

#[cfg(target_os = "linux")]
mod linux {
    use super::super::dl::Lib;
    use std::ffi::{c_int, c_void};

    /// The package zones' `energy_uj`, each wrapping at its range, summed.
    pub struct Rapl {
        zones: Vec<(String, u64, u64)>,
        total: u64,
    }

    fn read_u64(p: &str) -> Option<u64> {
        std::fs::read_to_string(p).ok()?.trim().parse().ok()
    }

    impl Rapl {
        pub fn open() -> Result<Rapl, String> {
            let mut zones = Vec::new();
            for e in std::fs::read_dir("/sys/class/powercap").map_err(|e| format!("no powercap zones: {e}"))?.flatten() {
                let p = e.path();
                let n = p.file_name().and_then(|n| n.to_str()).unwrap_or_default().to_string();
                if !(n.starts_with("intel-rapl:") && n.matches(':').count() == 1) || !std::fs::read_to_string(p.join("name")).is_ok_and(|s| s.trim().starts_with("package-")) {
                    continue;
                }
                let z = p.display().to_string();
                let range = read_u64(&format!("{z}/max_energy_range_uj")).ok_or("a package zone with no range")?.saturating_add(1);
                let now = read_u64(&format!("{z}/energy_uj")).ok_or_else(|| format!("{z}/energy_uj is not readable"))?;
                zones.push((z, range, now));
            }
            zones.sort();
            if zones.is_empty() { Err("no package zone".into()) } else { Ok(Rapl { zones, total: 0 }) }
        }
    }

    impl super::Counter for Rapl {
        fn name(&self) -> &'static str {
            "rapl"
        }
        fn attribution(&self) -> &'static str {
            "bracket"
        }
        fn fabrics(&self) -> Vec<&'static str> {
            vec!["cpu"]
        }
        fn read_nj(&mut self) -> Option<u64> {
            for z in &mut self.zones {
                let now = read_u64(&format!("{}/energy_uj", z.0))?;
                let d = if now >= z.2 { now - z.2 } else { z.1 - z.2 + now };
                if d > z.1 / 2 {
                    return None;
                }
                self.total += d * 1000;
                z.2 = now;
            }
            Some(self.total)
        }
    }

    /// The GPU driver's total-energy counter of device 0, in millijoules since the
    /// driver loaded, through its management library.
    pub struct Board {
        energy: unsafe extern "C" fn(*mut c_void, *mut u64) -> c_int,
        dev: *mut c_void,
    }

    impl Board {
        pub fn open() -> Result<Board, String> {
            let lib = Lib::open("libnvidia-ml.so.1")?; // allow-name: the GPU driver's management library file
            // SAFETY: each symbol cast to its declared signature; device 0's handle.
            unsafe {
                let init: unsafe extern "C" fn() -> c_int = lib.func("nvmlInit_v2")?; // allow-name: the library's entry point
                let by_index: unsafe extern "C" fn(u32, *mut *mut c_void) -> c_int = lib.func("nvmlDeviceGetHandleByIndex_v2")?; // allow-name: the library's entry point
                let energy = lib.func("nvmlDeviceGetTotalEnergyConsumption")?; // allow-name: the library's entry point
                if init() != 0 {
                    return Err("the management library did not initialise".into());
                }
                let mut dev = std::ptr::null_mut();
                if by_index(0, &mut dev) != 0 {
                    return Err("no device 0".into());
                }
                Ok(Board { energy, dev })
            }
        }
    }

    impl super::Counter for Board {
        fn name(&self) -> &'static str {
            "gpu-board"
        }
        fn attribution(&self) -> &'static str {
            "bracket"
        }
        fn fabrics(&self) -> Vec<&'static str> {
            vec!["gpu", "media"]
        }
        fn read_nj(&mut self) -> Option<u64> {
            let mut mj = 0u64;
            // SAFETY: the device handle the library returned.
            (unsafe { (self.energy)(self.dev, &mut mj) } == 0).then_some(mj * 1_000_000)
        }
    }

    pub struct Both(pub Rapl, pub Board);

    impl super::Counter for Both {
        fn name(&self) -> &'static str {
            "rapl-and-gpu-board"
        }
        fn attribution(&self) -> &'static str {
            "bracket"
        }
        fn fabrics(&self) -> Vec<&'static str> {
            vec!["cpu", "gpu", "media"]
        }
        fn read_nj(&mut self) -> Option<u64> {
            Some(self.0.read_nj()? + self.1.read_nj()?)
        }
    }
}

/// One round: its figure in microjoules (or why it has none), and its busy window.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Round {
    pub uj: Result<u64, String>,
    pub busy_us: u64,
    pub idle_us: u64,
}

/// Run `encode` once unmetered (its first run loads and warms the encoder), then
/// `rounds` times (odd, at least 3) around `counter`, a `bracket` counter's idle window
/// of `idle` before each, and seal the median round (§8.2). Returns the figure, every
/// metered round, and the last round's encode.
pub fn measure<T>(counter: &mut dyn Counter, rounds: usize, idle: Duration, mut encode: impl FnMut() -> Result<T, String>) -> Result<(Option<Energy>, Vec<Round>, T), String> {
    if rounds < 3 || rounds.is_multiple_of(2) {
        return Err("an odd number of rounds, at least 3".into());
    }
    encode()?;
    let bracket = counter.attribution() == "bracket";
    let mut out = Vec::with_capacity(rounds);
    let mut last = None;
    for _ in 0..rounds {
        let (mut idle_nj, mut idle_us) = (0u64, 0u64);
        if bracket {
            let (t0, e0) = (Instant::now(), counter.read_nj());
            std::thread::sleep(idle);
            let (t1, e1) = (Instant::now(), counter.read_nj());
            idle_us = (t1 - t0).as_micros() as u64;
            idle_nj = match (e0, e1) {
                (Some(a), Some(b)) => b - a,
                _ => u64::MAX,
            };
        }
        let (t0, e0) = (Instant::now(), counter.read_nj());
        let r = encode()?;
        let (t1, e1) = (Instant::now(), counter.read_nj());
        let busy_us = (t1 - t0).as_micros() as u64;
        let uj = match (e0, e1) {
            (Some(a), Some(b)) if idle_nj != u64::MAX => {
                let busy = b - a;
                // The idle window's mean power over the busy window's length.
                let base = if bracket && idle_us > 0 { (u128::from(idle_nj) * u128::from(busy_us) / u128::from(idle_us)) as u64 } else { 0 };
                if busy <= base {
                    Err("no marginal: the busy window drew no more than the idle one".to_string())
                } else {
                    Ok((busy - base).div_ceil(1000))
                }
            }
            _ => Err("a counter reading failed".to_string()),
        };
        out.push(Round { uj, busy_us, idle_us });
        last = Some(r);
    }
    let last = last.expect("at least one round");
    let mut figures: Vec<(u64, u64)> = Vec::new();
    for r in &out {
        match &r.uj {
            Ok(u) => figures.push((*u, r.busy_us)),
            Err(_) => return Ok((None, out, last)),
        }
    }
    figures.sort();
    let (median, window) = figures[figures.len() / 2];
    let (lo, hi) = (figures[0].0, figures[figures.len() - 1].0);
    if median == 0 || window == 0 {
        return Ok((None, out, last));
    }
    // The rounds' spread read as rectangular: (hi − lo) / (2·√3), relative to the
    // median, in ppm and rounded up (2·√3 < 3.4642); and the rounding floor of
    // energy-measurement §2.3.
    let spread_ppm = (u128::from(hi - lo) * 1_000_000 * 10_000).div_ceil(34_641 * u128::from(median)) as u64;
    let floor_ppm = 288_676u64.div_ceil(median);
    let relative_ppm = spread_ppm.max(floor_ppm).max(1).min(u64::from(u32::MAX));
    let e = Energy {
        joules_micro: median,
        class: "OnChipCounter",
        uncertainty: Some((relative_ppm, window)),
        filtering: Some("Unknown"),
        coverage: "Partial",
        counter: counter.name(),
        attribution: counter.attribution(),
        fabrics: counter.fabrics(),
    };
    Ok((Some(e), out, last))
}
