//! In-process workloads: c2c-latency (cache-line ping-pong between core
//! pairs) and dram-latency (pointer-chase through the memory hierarchy).
//! Zero-dep: CPU pinning goes through a raw `sched_setaffinity` syscall
//! (x86_64 nr 203); timing uses `std::time::Instant` (clock_gettime via
//! the vDSO — ns resolution, no TSC calibration needed).
//!
//! Both workloads measure while pinning threads, so they double as
//! coherency/memory stability checks: the ping-pong verifies every
//! cache-line handoff carries the exact expected value (a coherence or
//! clock failure shows up as a skipped/duplicated sequence number =
//! Verdict::Failed), and the pointer-chase verifies the chase walks one
//! full cycle of the randomized chain (a flipped address bit breaks the
//! cycle). Verdict is the stability signal; the latencies are the score.

use std::io;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use std::thread;
use std::time::{Duration, Instant};

use crate::workload::{C2cPair, DramLevel, Score, Verdict, Workload};

/// Parse a stress-ng `--taskset`-style cpu list ("0-11", "12-15",
/// "0-15", "3,7", ""/all = every logical cpu). Out-of-range entries are
/// clamped to the CPU count; invalid tokens are skipped. Sorted + deduped.
pub fn parse_cpu_list(cores: &str) -> Vec<usize> {
    let nproc = thread::available_parallelism()
        .map(|n| n.get())
        .unwrap_or(1);
    let trimmed = cores.trim();
    if trimmed.is_empty() || trimmed.eq_ignore_ascii_case("all") {
        return (0..nproc).collect();
    }
    let mut cpus = Vec::new();
    for token in trimmed.split(',') {
        let token = token.trim();
        if let Some((a, b)) = token.split_once('-') {
            if let (Ok(a), Ok(b)) = (a.trim().parse::<usize>(), b.trim().parse::<usize>()) {
                for c in a..=b {
                    if c < nproc {
                        cpus.push(c);
                    }
                }
            }
        } else if let Ok(c) = token.parse::<usize>() {
            if c < nproc {
                cpus.push(c);
            }
        }
    }
    cpus.sort_unstable();
    cpus.dedup();
    cpus
}

/// Pin the CALLING thread to one logical cpu. Raw syscall — std doesn't
/// expose sched_setaffinity and we don't want a libc dep. Returns false
/// if pinning isn't available (non-x86_64) or the kernel rejected it.
fn pin_self(cpu: usize) -> bool {
    let words = cpu / 64 + 1;
    let mut mask = vec![0u64; words];
    mask[cpu / 64] |= 1u64 << (cpu % 64);
    let ret =
        unsafe { sched_setaffinity_syscall(0, words * std::mem::size_of::<u64>(), mask.as_ptr()) };
    ret == 0
}

#[cfg(target_arch = "x86_64")]
unsafe fn sched_setaffinity_syscall(pid: i32, len: usize, mask: *const u64) -> isize {
    const SYS_SCHED_SETAFFINITY: isize = 203;
    let ret: isize;
    core::arch::asm!(
        "syscall",
        inlateout("rax") SYS_SCHED_SETAFFINITY => ret,
        in("rdi") pid,
        in("rsi") len,
        in("rdx") mask,
        lateout("rcx") _,
        lateout("r11") _,
        options(nostack)
    );
    ret
}

#[cfg(not(target_arch = "x86_64"))]
unsafe fn sched_setaffinity_syscall(_pid: i32, _len: usize, _mask: *const u64) -> isize {
    -1
}

// ---------------------------------------------------------------------------
// Shared in-process runner skeleton
// ---------------------------------------------------------------------------

/// Measurement body: runs to completion, honors `cancel` between units of
/// work, returns (verdict, score). `Stops early → (Stopped, None)`.
type Worker = fn(duration_secs: u64, cpus: &[usize], cancel: &AtomicBool) -> (Verdict, Score);

/// One struct for both in-process workloads: the trait plumbing (start /
/// is_running / stop / wait / score) is identical; only the worker fn
/// pointer + name differ.
struct InProc {
    name: &'static str,
    duration_secs: u64,
    cores: String,
    worker: Worker,
    handle: Option<thread::JoinHandle<()>>,
    cancel: Arc<AtomicBool>,
    done: Arc<AtomicBool>,
    result: Arc<Mutex<Option<(Verdict, Score)>>>,
}

impl InProc {
    fn new(name: &'static str, duration_secs: u64, cores: String, worker: Worker) -> Self {
        InProc {
            name,
            duration_secs,
            cores,
            worker,
            handle: None,
            cancel: Arc::new(AtomicBool::new(false)),
            done: Arc::new(AtomicBool::new(false)),
            result: Arc::new(Mutex::new(None)),
        }
    }
}

impl Workload for InProc {
    fn name(&self) -> &str {
        self.name
    }

    fn params(&self) -> String {
        format!("~{}s cores={}", self.duration_secs, self.cores)
    }

    fn start(&mut self) -> io::Result<()> {
        self.cancel.store(false, Ordering::Relaxed);
        self.done.store(false, Ordering::Relaxed);
        *self.result.lock().unwrap() = None;
        let duration = self.duration_secs;
        let cores = self.cores.clone();
        let cancel = Arc::clone(&self.cancel);
        let done = Arc::clone(&self.done);
        let result = Arc::clone(&self.result);
        let worker = self.worker;
        self.handle = Some(thread::spawn(move || {
            let cpus = parse_cpu_list(&cores);
            let out = worker(duration, &cpus, &cancel);
            *result.lock().unwrap() = Some(out);
            done.store(true, Ordering::Release);
        }));
        Ok(())
    }

    fn is_running(&mut self) -> bool {
        !self.done.load(Ordering::Acquire)
    }

    fn stop(&mut self) -> io::Result<()> {
        self.cancel.store(true, Ordering::Relaxed);
        Ok(())
    }

    fn wait(&mut self) -> Verdict {
        match self.handle.take() {
            Some(handle) => {
                if handle.join().is_err() {
                    return Verdict::Error("measurement thread panicked".to_string());
                }
                let guard = self.result.lock().unwrap();
                match &*guard {
                    Some((v, _)) => v.clone(),
                    None => Verdict::Stopped,
                }
            }
            None => Verdict::Stopped,
        }
    }

    fn score(&mut self) -> Score {
        let guard = self.result.lock().unwrap();
        match &*guard {
            Some((_, s)) => s.clone(),
            None => Score::None,
        }
    }
}

/// Public constructors used by workload::from_spec.
pub fn c2c_latency(duration_secs: u64, cores: String) -> Box<dyn Workload> {
    Box::new(InProc::new(
        "c2c-latency",
        duration_secs,
        cores,
        measure_c2c,
    ))
}

pub fn dram_latency(duration_secs: u64, cores: String) -> Box<dyn Workload> {
    Box::new(InProc::new(
        "dram-latency",
        duration_secs,
        cores,
        measure_dram,
    ))
}

// ---------------------------------------------------------------------------
// c2c-latency: cache-line ping-pong between pinned core pairs
// ---------------------------------------------------------------------------

/// Round trips per pair. ~2 ms of transfers per pair on this chip; the
/// full upper triangle of 16 logical cpus (120 pairs) lands well under
/// the 30 s default duration.
const C2C_ROUND_TRIPS: u64 = 20_000;

fn measure_c2c(_duration: u64, cpus: &[usize], cancel: &AtomicBool) -> (Verdict, Score) {
    if cpus.len() < 2 {
        return (
            Verdict::Error("c2c-latency needs at least 2 logical cpus".to_string()),
            Score::None,
        );
    }
    let mut pairs = Vec::new();
    for (ia, &a) in cpus.iter().enumerate() {
        for &b in &cpus[ia + 1..] {
            if cancel.load(Ordering::Relaxed) {
                return (Verdict::Stopped, Score::None);
            }
            match ping_pong(a, b) {
                Ok(ns_per_hop) => pairs.push(C2cPair {
                    cpu_a: a,
                    cpu_b: b,
                    ns_per_hop,
                }),
                Err(v) => return (v, Score::None),
            }
        }
    }
    (Verdict::Clean, Score::C2c { pairs })
}

/// Ping-pong a cache line (one AtomicU64) between cpu `a` and cpu `b`
/// for C2C_ROUND_TRIPS round trips. Each sequence value has exactly one
/// writer, so a plain spin-load that sees a value past its expectation
/// means the line was corrupted or the other side double-wrote — a
/// coherency/clock failure, reported as Verdict::Failed.
fn ping_pong(a: usize, b: usize) -> Result<f64, Verdict> {
    let ball = Arc::new(AtomicU64::new(0));
    let corrupt = Arc::new(AtomicBool::new(false));
    let pin_fail = Arc::new(AtomicBool::new(false));

    let ball_a = Arc::clone(&ball);
    let ball_b = Arc::clone(&ball);
    let corrupt_a = Arc::clone(&corrupt);
    let corrupt_b = Arc::clone(&corrupt);
    let pf_a = Arc::clone(&pin_fail);
    let pf_b = Arc::clone(&pin_fail);

    // Initiator: A. Stores 1, waits for 2, stores 3, … timing brackets
    // the full 0 → 2N transfer sequence.
    let tx = thread::spawn(move || {
        if !pin_self(a) {
            pf_a.store(true, Ordering::Relaxed);
        }
        let t0 = Instant::now();
        for i in 0..C2C_ROUND_TRIPS {
            let expect = 2 * i;
            loop {
                let v = ball_a.load(Ordering::Relaxed);
                if v == expect {
                    break;
                }
                if v > expect {
                    corrupt_a.store(true, Ordering::Relaxed);
                    return Duration::ZERO;
                }
                std::hint::spin_loop();
            }
            ball_a.store(expect + 1, Ordering::Relaxed);
        }
        // Wait for B's final store so the timing covers 2N transfers.
        let final_expect = 2 * C2C_ROUND_TRIPS;
        loop {
            let v = ball_a.load(Ordering::Relaxed);
            if v == final_expect {
                break;
            }
            if v > final_expect {
                corrupt_a.store(true, Ordering::Relaxed);
                return Duration::ZERO;
            }
            std::hint::spin_loop();
        }
        t0.elapsed()
    });
    // Responder: B. Waits for 1, stores 2, waits for 3, …
    let rx = thread::spawn(move || {
        if !pin_self(b) {
            pf_b.store(true, Ordering::Relaxed);
        }
        for i in 0..C2C_ROUND_TRIPS {
            let expect = 2 * i + 1;
            loop {
                let v = ball_b.load(Ordering::Relaxed);
                if v == expect {
                    break;
                }
                if v > expect {
                    corrupt_b.store(true, Ordering::Relaxed);
                    return;
                }
                std::hint::spin_loop();
            }
            ball_b.store(expect + 1, Ordering::Relaxed);
        }
    });

    let elapsed = match (tx.join(), rx.join()) {
        (Ok(e), Ok(())) => e,
        _ => return Err(Verdict::Error("ping-pong thread panicked".to_string())),
    };
    if pin_fail.load(Ordering::Relaxed) {
        return Err(Verdict::Error(format!(
            "can't pin cpus {a}/{b} (sched_setaffinity failed)"
        )));
    }
    if corrupt.load(Ordering::Relaxed) {
        return Err(Verdict::Failed(format!(
            "cache-line handoff {a}<->{b} corrupted"
        )));
    }
    let transfers = 2.0 * C2C_ROUND_TRIPS as f64;
    Ok(elapsed.as_nanos() as f64 / transfers)
}

// ---------------------------------------------------------------------------
// dram-latency: pointer-chase through the cache hierarchy
// ---------------------------------------------------------------------------

/// Chase buffers, bytes. Sizes straddle the 12600KF hierarchy: L1D 48K
/// per P-core, L2 2MB per P-core, L3 20MB shared. 32K fits L1, 512K fits
/// L2, 8M fits L3, 64M falls through to DRAM.
const DRAM_SIZES: [u64; 4] = [32 * 1024, 512 * 1024, 8 * 1024 * 1024, 64 * 1024 * 1024];

/// Minimum time per size so the slow levels (DRAM) get enough samples.
const DRAM_MIN_PER_SIZE: Duration = Duration::from_millis(200);

fn measure_dram(_duration: u64, cpus: &[usize], cancel: &AtomicBool) -> (Verdict, Score) {
    let Some(&cpu) = cpus.first() else {
        return (
            Verdict::Error("dram-latency needs at least 1 cpu".to_string()),
            Score::None,
        );
    };
    // Runs on the InProc worker thread itself — pin it (the thread exits
    // with the measurement, so no affinity restore needed).
    if !pin_self(cpu) {
        return (
            Verdict::Error(format!("can't pin cpu {cpu} (sched_setaffinity failed)")),
            Score::None,
        );
    }
    let mut levels = Vec::new();
    for &size in &DRAM_SIZES {
        if cancel.load(Ordering::Relaxed) {
            return (Verdict::Stopped, Score::None);
        }
        let lines = (size / 64) as usize;
        let table = build_chase_table(lines);
        match chase(&table, lines, cancel) {
            Ok(ns) if ns.is_nan() => return (Verdict::Stopped, Score::None), // cancelled mid-size
            Ok(ns) => levels.push(DramLevel {
                size_kb: size / 1024,
                ns_per_access: ns,
            }),
            Err(v) => return (v, Score::None),
        }
    }
    (Verdict::Clean, Score::Dram { levels })
}

/// Random single cycle over `lines` cache lines, laid out in a u64 table
/// with one u64 payload per 64-byte line. `table[line*8]` = next line.
/// A cycle of length exactly `lines` starting at line 0.
fn build_chase_table(lines: usize) -> Vec<u64> {
    let mut order: Vec<u32> = (0..lines as u32).collect();
    // Fisher-Yates with a xorshift RNG (deterministic per run — fine for
    // a latency measurement; the chase only needs to be irregular).
    let mut state: u64 = 0x9E3779B97F4A7C15 ^ lines as u64;
    for i in (1..lines).rev() {
        state ^= state << 13;
        state ^= state >> 7;
        state ^= state << 17;
        let j = (state % (i as u64 + 1)) as usize;
        order.swap(i, j);
    }
    // order[0] is the head (chase start). link[order[i]] = order[i+1].
    let mut table = vec![0u64; lines * 8];
    for i in 0..lines {
        let line = order[i] as usize;
        table[line * 8] = order[(i + 1) % lines] as u64;
    }
    table
}

/// Chase full laps of the cycle until DRAM_MIN_PER_SIZE elapsed (or
/// cancel). Returns ns/access. The table is ONE cycle covering all
/// `lines` nodes, so starting anywhere, exactly `lines` dependent loads
/// must land back on the start — if a stored address is ever wrong
/// (flipped bit), the lap check fails: that's the memory-stability
/// self-check.
fn chase(table: &[u64], lines: usize, cancel: &AtomicBool) -> Result<f64, Verdict> {
    let mut idx = 0u64;
    let mut laps = 0u64;
    let t0 = Instant::now();
    loop {
        for _ in 0..lines {
            idx = std::hint::black_box(table[idx as usize * 8]);
        }
        laps += 1;
        if idx != 0 {
            return Err(Verdict::Failed(
                "pointer-chase left the cycle (memory fault?)".to_string(),
            ));
        }
        if t0.elapsed() >= DRAM_MIN_PER_SIZE || laps > 100_000_000 {
            break;
        }
        if cancel.load(Ordering::Relaxed) {
            return Ok(f64::NAN); // ignored — caller checks cancel separately
        }
    }
    let accesses = laps as f64 * lines as f64;
    Ok(t0.elapsed().as_nanos() as f64 / accesses)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn parse_cpu_list_forms() {
        let n = thread::available_parallelism()
            .map(|n| n.get())
            .unwrap_or(1);
        assert_eq!(parse_cpu_list("").len(), n);
        assert_eq!(parse_cpu_list("all").len(), n);
        assert_eq!(parse_cpu_list("2-4"), vec![2, 3, 4]);
        assert_eq!(parse_cpu_list("3,1,2-2"), vec![1, 2, 3]); // sort + dedup
        assert_eq!(parse_cpu_list("junk"), Vec::<usize>::new());
        // Out-of-range clamped away (cpu id >= nproc dropped).
        assert_eq!(parse_cpu_list("0-999").len(), n);
    }

    #[test]
    fn chase_table_is_one_cycle() {
        let lines = 4096;
        let table = build_chase_table(lines);
        let mut idx = 0u64;
        for _ in 0..lines {
            idx = table[idx as usize * 8];
            assert!(idx < lines as u64);
        }
        assert_eq!(
            idx, 0,
            "single cycle must return to the start after `lines` hops"
        );
    }

    #[test]
    fn chase_small_table_measures_ns() {
        let lines = 1024;
        let table = build_chase_table(lines);
        let cancel = AtomicBool::new(false);
        let ns = chase(&table, lines, &cancel).unwrap();
        // Only a floor + deadlock cap: absolute ns bounds are meaningless
        // when the suite runs on a box under load (the dev machine's
        // benchmarks), and the hierarchy assertions live in dram_e2e.
        assert!(ns > 0.0 && ns < 1e6, "sanity range, got {ns}");
    }

    #[test]
    fn chase_fails_on_desynced_lap() {
        let lines = 1024;
        let mut table = build_chase_table(lines);
        // Rewire one edge to skip a node: the cycle splits, so a
        // `lines`-hop lap from 0 no longer lands back on 0.
        let a = table[0];
        let b = table[a as usize * 8];
        table[a as usize * 8] = table[b as usize * 8];
        let cancel = AtomicBool::new(false);
        let r = chase(&table, lines, &cancel);
        assert!(matches!(r, Err(Verdict::Failed(_))), "got {r:?}");
    }

    #[test]
    #[cfg_attr(not(target_arch = "x86_64"), ignore)] // pinning is x86_64-only here
    fn ping_pong_between_two_cpus() {
        let cpus = parse_cpu_list("0-15");
        if cpus.len() < 2 {
            return; // single-cpu box (container?) — nothing to test
        }
        let ns = ping_pong(cpus[0], cpus[1]).expect("clean handoff");
        // Floor + deadlock cap only — a preempted responder (this test
        // runs on a box that may be mid-benchmark) legitimately pushes
        // handoffs way past the tens-of-ns idle range.
        assert!(ns > 0.0 && ns < 1e7, "sanity range, got {ns}");
    }

    #[test]
    #[cfg_attr(not(target_arch = "x86_64"), ignore)]
    fn c2c_e2e_via_trait() {
        use crate::config::WorkloadSpec;
        let cpus = parse_cpu_list("0-15");
        if cpus.len() < 2 {
            return;
        }
        let spec = WorkloadSpec {
            name: "c2c-latency".to_string(),
            duration_secs: 1,
            cores: "0-3".to_string(),
        };
        let mut w = crate::workload::from_spec(&spec, std::path::Path::new("bin")).unwrap();
        assert_eq!(w.name(), "c2c-latency");
        w.start().unwrap();
        while w.is_running() {
            thread::sleep(Duration::from_millis(25));
        }
        assert_eq!(w.wait(), Verdict::Clean);
        match w.score() {
            Score::C2c { pairs } => {
                // 4 cpus → 6 upper-triangle pairs.
                assert_eq!(pairs.len(), 6);
                assert!(pairs.iter().all(|p| p.ns_per_hop > 0.0));
            }
            other => panic!("expected C2c score, got {other:?}"),
        }
    }

    #[test]
    #[cfg_attr(not(target_arch = "x86_64"), ignore)]
    fn dram_e2e_via_trait() {
        use crate::config::WorkloadSpec;
        let spec = WorkloadSpec {
            name: "dram-latency".to_string(),
            duration_secs: 1,
            cores: "0".to_string(),
        };
        let mut w = crate::workload::from_spec(&spec, std::path::Path::new("bin")).unwrap();
        w.start().unwrap();
        while w.is_running() {
            thread::sleep(Duration::from_millis(50));
        }
        assert_eq!(w.wait(), Verdict::Clean);
        match w.score() {
            Score::Dram { levels } => {
                assert_eq!(levels.len(), 4);
                // DRAM (64 MiB) must be far slower than the L1-sized set.
                assert!(levels[0].ns_per_access < levels[3].ns_per_access);
                assert!(
                    levels[3].ns_per_access > 30.0,
                    "DRAM tier too fast: {levels:?}"
                );
            }
            other => panic!("expected Dram score, got {other:?}"),
        }
    }

    #[test]
    fn stop_mid_run_yields_stopped() {
        let cpus = parse_cpu_list("0-15");
        if cpus.len() < 2 {
            return;
        }
        let mut w = c2c_latency(3600, "0-15".to_string());
        w.start().unwrap();
        w.stop().unwrap();
        while w.is_running() {
            thread::sleep(Duration::from_millis(25));
        }
        assert_eq!(w.wait(), Verdict::Stopped);
    }
}