//! Concurrent-read harness: how does decoded read throughput scale with the
//! number of threads reading one open file?
//!
//! libhdf5 (threadsafe build) serialises every API call under one global
//! mutex, and h5py holds it too, so threads cannot decode in parallel there.
//! A clawhdf5 [`File`] is `Send + Sync`; this harness measures what that buys.
//! `crates/clawhdf5-bench/scripts/concurrent_read_h5py.py` runs the same
//! workload on the same files with h5py (threads, and processes), and
//! `compare_concurrent_read.py` tabulates the JSON both write.
//!
//! Files (generated on first use, reused while `manifest.json` matches):
//!
//! * `<dir>/deflate.h5`: `--datasets` datasets `d00`, `d01`, ... of `f32`,
//!   `--mib` MiB decoded each, shape `[mib * 256, 1024]`, chunks `256 x 256`,
//!   deflate level 4.
//! * `<dir>/contiguous.h5`: the same datasets, contiguous.
//!
//! Modes, for each layout and each thread count `T` (strong scaling: the total
//! work per repetition is fixed, split among the threads):
//!
//! * `distinct`: every dataset is read in full once; thread `t` reads datasets
//!   `t, t + T, t + 2T, ...`.
//! * `same`: all threads read `d00`, `--slabs` random `--slab` x `--slab`
//!   hyperslabs in total (slab `j` goes to thread `j % T`). The offsets come
//!   from a splitmix64 stream seeded with `--seed`, identical in the h5py
//!   script.
//!
//! One `File` per layout per repetition is shared by all threads (opened
//! fresh each repetition, so no chunk cache carries over). Page cache:
//! `warm` (default) reads every file once before timing; `--cold` evicts the
//! files from the page cache with `posix_fadvise(POSIX_FADV_DONTNEED)` before
//! every repetition (no root needed; it only evicts clean, unmapped pages, so
//! it is best effort — the JSON says which was used).
//!
//! Decode inside one read is itself parallel when clawhdf5-format's `parallel`
//! feature is on (it is in this binary, via clawhdf5-agent). `--decode-threads
//! N` sizes that rayon pool; `--decode-threads 1` measures the API's own
//! thread scaling, comparable with h5py where each call decodes on the
//! calling thread.
//!
//! ```text
//! cargo run --release -p clawhdf5-bench --bin concurrent_read -- \
//!     --dir /data/concurrent-read --json clawhdf5.json
//! cargo run --release -p clawhdf5-bench --bin concurrent_read -- \
//!     --dir /tmp/cr --datasets 4 --mib 1 --threads 1,2 --slabs 16 --reps 1   # smoke
//! ```

use std::path::{Path, PathBuf};
use std::sync::Barrier;
use std::time::Instant;

use clawhdf5::{File, FileBuilder, Selection};
use serde::{Deserialize, Serialize};

const COLS: u64 = 1024;
const ROWS_PER_MIB: u64 = 256; // 256 rows x 1024 cols x 4 bytes = 1 MiB
const CHUNK: u64 = 256;
const DEFLATE_LEVEL: u32 = 4;
const LAYOUTS: [&str; 2] = ["deflate", "contiguous"];
const MANIFEST_VERSION: u32 = 1;

/// splitmix64 — shared with the h5py script, which must produce the same
/// stream (both the data and the hyperslab offsets depend on it).
fn splitmix64(state: &mut u64) -> u64 {
    *state = state.wrapping_add(0x9E37_79B9_7F4A_7C15);
    let mut z = *state;
    z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
    z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
    z ^ (z >> 31)
}

/// Element `i` (row-major) of dataset `k`: a slowly varying integer part plus
/// 8 bits of noise, so deflate has real work to do (about 3.1x) and every value
/// is exact in `f32` (< 2^15 with 8 fraction bits), which lets both harnesses
/// check what they read against this formula.
fn value(k: u64, i: u64) -> f32 {
    let mut s = i ^ (k << 40);
    let noise = splitmix64(&mut s) & 0xff;
    (((i >> 6) % 16384) + k) as f32 + noise as f32 / 256.0
}

#[derive(Serialize, Deserialize, PartialEq, Debug, Clone)]
struct Manifest {
    version: u32,
    datasets: u64,
    rows: u64,
    cols: u64,
    chunk: [u64; 2],
    deflate_level: u32,
    files: Vec<(String, String)>, // (layout, file name)
    writer: String,
}

fn manifest_for(datasets: u64, mib: u64) -> Manifest {
    Manifest {
        version: MANIFEST_VERSION,
        datasets,
        rows: mib * ROWS_PER_MIB,
        cols: COLS,
        chunk: [CHUNK, CHUNK],
        deflate_level: DEFLATE_LEVEL,
        files: LAYOUTS
            .iter()
            .map(|l| (l.to_string(), format!("{l}.h5")))
            .collect(),
        writer: format!("clawhdf5 {}", env!("CARGO_PKG_VERSION")),
    }
}

fn dataset_values(k: u64, n: u64) -> Vec<f32> {
    (0..n).map(|i| value(k, i)).collect()
}

/// Write the files unless `dir` already holds ones matching `want`.
fn ensure_files(dir: &Path, want: &Manifest) -> std::io::Result<bool> {
    let manifest_path = dir.join("manifest.json");
    if let Ok(text) = std::fs::read_to_string(&manifest_path)
        && let Ok(have) = serde_json::from_str::<Manifest>(&text)
        && have.version == want.version
        && have.datasets == want.datasets
        && have.rows == want.rows
        && have.cols == want.cols
        && have.chunk == want.chunk
        && have.deflate_level == want.deflate_level
        && have.files == want.files
        && want.files.iter().all(|(_, f)| dir.join(f).exists())
    {
        return Ok(false);
    }
    std::fs::create_dir_all(dir)?;
    // A stale manifest must not survive a half-written regeneration.
    let _ = std::fs::remove_file(&manifest_path);
    let n = want.rows * want.cols;
    for (layout, file) in &want.files {
        // One layout at a time keeps the peak memory to about twice one
        // file's decoded size.
        let mut b = FileBuilder::new();
        for k in 0..want.datasets {
            let ds = b.create_dataset(&format!("d{k:02}"));
            ds.with_f32_data(&dataset_values(k, n))
                .with_shape(&[want.rows, want.cols]);
            if layout == "deflate" {
                ds.with_chunks(&[CHUNK.min(want.rows), CHUNK])
                    .with_deflate(DEFLATE_LEVEL);
            }
        }
        b.write(dir.join(file)).map_err(std::io::Error::other)?;
    }
    std::fs::write(
        &manifest_path,
        serde_json::to_string_pretty(want).map_err(std::io::Error::other)?,
    )?;
    Ok(true)
}

fn slab_offsets(seed: u64, count: usize, rows: u64, cols: u64, slab: u64) -> Vec<(u64, u64)> {
    let mut s = seed;
    (0..count)
        .map(|_| {
            let r = splitmix64(&mut s) % (rows - slab + 1);
            let c = splitmix64(&mut s) % (cols - slab + 1);
            (r, c)
        })
        .collect()
}

/// Warm the page cache by reading every byte of `path`.
fn warm(path: &Path) -> std::io::Result<()> {
    let mut f = std::fs::File::open(path)?;
    std::io::copy(&mut f, &mut std::io::sink())?;
    Ok(())
}

/// Ask the kernel to drop `path`'s pages from the page cache.
fn evict(path: &Path) -> std::io::Result<()> {
    use std::os::fd::AsRawFd;
    let f = std::fs::File::open(path)?;
    // SAFETY: plain syscall on a valid, open file descriptor.
    let rc = unsafe { libc::posix_fadvise(f.as_raw_fd(), 0, 0, libc::POSIX_FADV_DONTNEED) };
    if rc != 0 {
        return Err(std::io::Error::from_raw_os_error(rc));
    }
    Ok(())
}

#[derive(Serialize)]
struct Row {
    layout: String,
    mode: String,
    threads: usize,
    /// Decoded (selected) bytes read per repetition.
    bytes: u64,
    times_s: Vec<f64>,
    median_s: f64,
    mb_s: f64,
    /// `mb_s / (threads * mb_s at threads = 1)`; null without a 1-thread row.
    efficiency: Option<f64>,
}

struct Args {
    dir: PathBuf,
    datasets: u64,
    mib: u64,
    threads: Vec<usize>,
    reps: usize,
    slab: u64,
    slabs: usize,
    seed: u64,
    cold: bool,
    decode_threads: usize,
    modes: Vec<String>,
    layouts: Vec<String>,
    json: Option<PathBuf>,
}

const USAGE: &str = "\
usage: concurrent_read [--dir DIR] [--datasets N] [--mib N] [--threads 1,2,4,8,16]
                       [--reps N] [--slab N] [--slabs N] [--seed N] [--cold]
                       [--decode-threads N] [--modes distinct,same]
                       [--layouts deflate,contiguous] [--json FILE]";

fn parse_list<T: std::str::FromStr>(s: &str) -> Result<Vec<T>, String> {
    s.split(',')
        .map(|x| x.trim().parse().map_err(|_| format!("bad list item {x:?}")))
        .collect()
}

fn parse_args() -> Result<Args, String> {
    let mut a = Args {
        dir: PathBuf::from("concurrent-read-data"),
        datasets: 64,
        mib: 64,
        threads: vec![1, 2, 4, 8, 16],
        reps: 3,
        slab: 256,
        slabs: 1024,
        seed: 42,
        cold: false,
        decode_threads: 0,
        modes: vec!["distinct".into(), "same".into()],
        layouts: LAYOUTS.iter().map(|s| s.to_string()).collect(),
        json: None,
    };
    let mut it = std::env::args().skip(1);
    while let Some(flag) = it.next() {
        if flag == "--cold" {
            a.cold = true;
            continue;
        }
        if flag == "-h" || flag == "--help" {
            return Err(USAGE.into());
        }
        let v = it.next().ok_or(format!("{flag} needs a value\n{USAGE}"))?;
        let num = |v: &str| {
            v.parse::<u64>()
                .map_err(|_| format!("{flag}: bad number {v:?}"))
        };
        match flag.as_str() {
            "--dir" => a.dir = v.into(),
            "--datasets" => a.datasets = num(&v)?,
            "--mib" => a.mib = num(&v)?,
            "--threads" => a.threads = parse_list(&v)?,
            "--reps" => a.reps = num(&v)? as usize,
            "--slab" => a.slab = num(&v)?,
            "--slabs" => a.slabs = num(&v)? as usize,
            "--seed" => a.seed = num(&v)?,
            "--decode-threads" => a.decode_threads = num(&v)? as usize,
            "--modes" => a.modes = parse_list(&v)?,
            "--layouts" => a.layouts = parse_list(&v)?,
            "--json" => a.json = Some(v.into()),
            _ => return Err(format!("unknown flag {flag}\n{USAGE}")),
        }
    }
    if a.datasets == 0 || a.datasets > 100 {
        return Err("--datasets must be 1..=100".into());
    }
    if a.mib == 0 || a.reps == 0 || a.slabs == 0 || a.threads.contains(&0) {
        return Err("--mib, --reps, --slabs and every --threads value must be > 0".into());
    }
    if a.slab == 0 || a.slab > COLS || a.slab > a.mib * ROWS_PER_MIB {
        return Err(format!(
            "--slab must be 1..={}",
            COLS.min(a.mib * ROWS_PER_MIB)
        ));
    }
    for m in &a.modes {
        if m != "distinct" && m != "same" {
            return Err(format!("unknown mode {m:?}"));
        }
    }
    for l in &a.layouts {
        if !LAYOUTS.contains(&l.as_str()) {
            return Err(format!("unknown layout {l:?}"));
        }
    }
    Ok(a)
}

/// One timed repetition: `T` threads on one shared `File`. Returns seconds.
fn run_once(
    path: &Path,
    mode: &str,
    threads: usize,
    m: &Manifest,
    slabs: &[(u64, u64)],
    slab: u64,
    verify: bool,
) -> f64 {
    let file = File::open(path).expect("open");
    let barrier = Barrier::new(threads + 1); // + the spawning thread
    let n = m.rows * m.cols;
    // Each thread times itself from the barrier; the repetition spans the
    // earliest start to the latest finish (timing on the spawning thread
    // instead undercounts whenever it is scheduled after the workers ran).
    let spans: Vec<(Instant, Instant)> = std::thread::scope(|s| {
        let handles: Vec<_> = (0..threads)
            .map(|t| {
                let (file, barrier) = (&file, &barrier);
                s.spawn(move || {
                    barrier.wait();
                    let start = Instant::now();
                    match mode {
                        "distinct" => {
                            for k in (t as u64..m.datasets).step_by(threads) {
                                let got = file.dataset(&format!("d{k:02}")).unwrap().read_f32();
                                let got = got.unwrap();
                                assert_eq!(got.len() as u64, n);
                                if verify {
                                    for i in [0, n / 3, n - 1] {
                                        assert_eq!(got[i as usize], value(k, i), "d{k:02}[{i}]");
                                    }
                                }
                                std::hint::black_box(got);
                            }
                        }
                        _ => {
                            let ds = file.dataset("d00").unwrap();
                            for &(r, c) in slabs.iter().skip(t).step_by(threads) {
                                let sel = Selection::Hyperslab {
                                    start: vec![r, c],
                                    stride: vec![1, 1],
                                    count: vec![slab, slab],
                                    block: vec![1, 1],
                                };
                                let got = ds.read_f32_selection(&sel).unwrap();
                                assert_eq!(got.len() as u64, slab * slab);
                                if verify {
                                    let last = (r + slab - 1) * m.cols + c + slab - 1;
                                    assert_eq!(got[0], value(0, r * m.cols + c));
                                    assert_eq!(*got.last().unwrap(), value(0, last));
                                }
                                std::hint::black_box(got);
                            }
                        }
                    }
                    (start, Instant::now())
                })
            })
            .collect();
        barrier.wait();
        handles.into_iter().map(|h| h.join().unwrap()).collect()
    });
    let start = spans.iter().map(|s| s.0).min().unwrap();
    let end = spans.iter().map(|s| s.1).max().unwrap();
    (end - start).as_secs_f64()
}

fn median(v: &[f64]) -> f64 {
    let mut s = v.to_vec();
    s.sort_by(f64::total_cmp);
    s[s.len() / 2]
}

fn hostname() -> String {
    std::fs::read_to_string("/proc/sys/kernel/hostname")
        .map(|s| s.trim().to_string())
        .unwrap_or_else(|_| "unknown".into())
}

fn main() {
    let args = match parse_args() {
        Ok(a) => a,
        Err(e) => {
            eprintln!("{e}");
            std::process::exit(2);
        }
    };
    if cfg!(debug_assertions) {
        eprintln!("warning: debug build — numbers are meaningless. Use --release.");
    }
    if args.decode_threads > 0 {
        rayon::ThreadPoolBuilder::new()
            .num_threads(args.decode_threads)
            .build_global()
            .expect("configure rayon pool");
    }

    let manifest = manifest_for(args.datasets, args.mib);
    let t = Instant::now();
    match ensure_files(&args.dir, &manifest) {
        Ok(true) => eprintln!(
            "generated {} in {:.1} s",
            args.dir.display(),
            t.elapsed().as_secs_f64()
        ),
        Ok(false) => eprintln!("reusing {}", args.dir.display()),
        Err(e) => {
            eprintln!("cannot write test files in {}: {e}", args.dir.display());
            std::process::exit(1);
        }
    }
    let path_of = |layout: &str| args.dir.join(format!("{layout}.h5"));
    let slabs = slab_offsets(
        args.seed,
        args.slabs,
        manifest.rows,
        manifest.cols,
        args.slab,
    );
    let dataset_bytes = manifest.rows * manifest.cols * 4;

    let mut rows: Vec<Row> = Vec::new();
    println!("| layout | mode | threads | MB/s | efficiency | median s |");
    println!("|---|---|---:|---:|---:|---:|");
    for layout in &args.layouts {
        let path = path_of(layout);
        // Untimed pass: page cache warm (unless --cold), results checked.
        if !args.cold {
            warm(&path).expect("warm page cache");
        }
        for mode in &args.modes {
            run_once(&path, mode, 1, &manifest, &slabs, args.slab, true);
            let bytes = match mode.as_str() {
                "distinct" => dataset_bytes * manifest.datasets,
                _ => args.slab * args.slab * 4 * args.slabs as u64,
            };
            let mut base: Option<f64> = None;
            for &threads in &args.threads {
                let times: Vec<f64> = (0..args.reps)
                    .map(|_| {
                        if args.cold {
                            evict(&path).expect("posix_fadvise");
                        }
                        run_once(&path, mode, threads, &manifest, &slabs, args.slab, false)
                    })
                    .collect();
                let med = median(&times);
                let mb_s = bytes as f64 / (1 << 20) as f64 / med;
                if threads == 1 {
                    base = Some(mb_s);
                }
                let efficiency = base.map(|b| mb_s / (threads as f64 * b));
                println!(
                    "| {layout} | {mode} | {threads} | {mb_s:.0} | {} | {med:.4} |",
                    efficiency.map_or("-".into(), |e| format!("{e:.2}"))
                );
                rows.push(Row {
                    layout: layout.clone(),
                    mode: mode.clone(),
                    threads,
                    bytes,
                    times_s: times,
                    median_s: med,
                    mb_s,
                    efficiency,
                });
            }
        }
    }

    if let Some(out) = &args.json {
        let doc = serde_json::json!({
            "tool": "clawhdf5",
            "version": env!("CARGO_PKG_VERSION"),
            "host": hostname(),
            "cpus": std::thread::available_parallelism().map_or(0, |n| n.get()),
            "unix_time": std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .map_or(0, |d| d.as_secs()),
            "cache": if args.cold { "cold (posix_fadvise DONTNEED before each repetition)" } else { "warm" },
            "decode_threads": rayon::current_num_threads(),
            "params": {
                "datasets": manifest.datasets,
                "mib": args.mib,
                "rows": manifest.rows,
                "cols": manifest.cols,
                "chunk": manifest.chunk,
                "deflate_level": manifest.deflate_level,
                "slab": args.slab,
                "slabs": args.slabs,
                "seed": args.seed,
                "reps": args.reps,
                "dir": args.dir,
            },
            "results": rows,
        });
        std::fs::write(out, serde_json::to_string_pretty(&doc).unwrap()).expect("write json");
    }
}

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

    #[test]
    fn values_are_exact_in_f32() {
        for k in [0, 7, 63] {
            for i in [0u64, 1, 4095, 1 << 20, (1 << 24) - 1] {
                let v = value(k, i);
                assert_eq!(v, (v as f64) as f32);
                assert!(v < 32768.0);
                assert_eq!((v * 256.0).fract(), 0.0);
            }
        }
    }

    /// The h5py script hard-codes this vector to check its splitmix64 port.
    #[test]
    fn splitmix64_reference() {
        let mut s = 42;
        assert_eq!(splitmix64(&mut s), 0xBDD7_3226_2FEB_6E95);
    }
}