i0 0.1.5

use clap::Parser;
use i0::{Timestamp, v3};
use rand::RngCore;
use std::num::NonZero;
use std::process::exit;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
use std::thread::{self, JoinHandle};
use std::time::{Duration, Instant};
use sysinfo::System;

const TEST_SEED: [u8; 32] = [
    157, 97, 177, 157, 239, 253, 90, 96, 186, 132, 74, 244, 146, 236, 44, 196, 68, 73, 197, 105,
    123, 50, 105, 25, 112, 59, 172, 3, 28, 174, 127, 96,
];

const NAME: &str = env!("CARGO_PKG_NAME");
const VERSION: &str = env!("CARGO_PKG_VERSION");

#[derive(Parser, Debug)]
#[command(about = format!("Tests how fast your device can generate thei.rs/{NAME}:3\n(v{VERSION})"), long_about = None)]
struct Cli {
    /// number of threads to use for searching
    #[arg(short, long)]
    nproc: Option<NonZero<usize>>,

    /// whether to seed the RNG – this might eliminate slowdowns caused by the RNG, but increases the chance for unintentional optimization falsifying the results
    #[arg(short, long, default_value_t = false)]
    seeded: bool,

    /// how long to dry-run the test before measuring throughput
    #[arg(short, long, default_value = "2s")]
    warmup: humantime::Duration,

    /// how long to run the test for (excluding warmup time)
    #[arg(short, long, default_value = "5s")]
    time: humantime::Duration,
}

fn main() {
    let cli = Cli::parse();
    let nproc = cli
        .nproc
        .unwrap_or_else(|| thread::available_parallelism().expect("unable to determine parallelism, manually specify the number of threads to run using -n"))
        .get();

    let t_warm = *cli.warmup;
    let t_run = *cli.time;
    if t_run.is_zero() {
        eprintln!("don't be silly, you can't run a benchmark with a 0s duration!");
        exit(1)
    }

    print!("{NAME}.rs[v{VERSION}]");
    if sysinfo::IS_SUPPORTED_SYSTEM {
        let mut sys = System::new_all();
        sys.refresh_all();

        let host_fqdn = System::host_name().unwrap_or("<unknown>".to_string());
        let host = host_fqdn.split(".").nth(0).unwrap();

        let cpu = if let Some(cpu) = sys.cpus().first() {
            cpu.brand()
        } else {
            "unknown cpu"
        };
        println!(" on {cpu} ({host})");
    } else {
        println!();
    }

    if nproc != 1 {
        run(cli.seeded, 1, t_warm, t_run);
    }

    run(cli.seeded, nproc, t_warm, t_run);
}

fn run(seeded: bool, threads: usize, t_warm: Duration, t_run: Duration) {
    let total = Arc::new(AtomicU64::new(0));
    let interrupt = Arc::new(AtomicBool::new(false));
    let warmup = Arc::new(AtomicBool::new(true));
    let active_worker_count = Arc::new(AtomicUsize::new(0));
    let mut handles = vec![];

    for _ in 0..threads {
        let total = Arc::clone(&total);
        let interrupt = Arc::clone(&interrupt);
        let warmup = Arc::clone(&warmup);
        let active_worker_count = Arc::clone(&active_worker_count);

        let runner = if seeded { run_seeded } else { run_unseeded };
        let handle = runner(total, interrupt, warmup, active_worker_count);

        handles.push(handle);
    }

    let handle = thread::spawn(move || {
        let prefix = format!(
            "{:>8}[{threads:>3}] ",
            if seeded { "seeded" } else { "unseeded" }
        );
        // warmup
        while active_worker_count.load(Ordering::Relaxed) < threads {
            thread::sleep(Duration::from_millis(1));
        }
        if !t_warm.is_zero() {
            println!("{prefix} warming up ({t_warm:?})");
            thread::sleep(t_warm);
            print!("\x1b[F");
        }
        // start
        println!("{prefix} generating i0:3 ({t_run:?})");
        warmup.store(false, Ordering::Relaxed);
        let start = Instant::now();
        thread::sleep(t_run);
        // end
        interrupt.store(true, Ordering::Relaxed);
        let elapsed = start.elapsed();
        while active_worker_count.load(Ordering::Relaxed) > 0 {
            thread::sleep(Duration::from_millis(1));
        }
        // result
        let ops = total.load(Ordering::Relaxed);
        let throughput = u128::from(ops) * 1_000_000_000 / elapsed.as_nanos();
        print!("\x1b[F");
        println!(
            "{prefix} {:>10} ops/sec {}",
            print_with_sep(throughput),
            " ".repeat(10)
        );
    });
    handle.join().unwrap();
    for handle in handles {
        handle.join().unwrap();
    }
}

fn run_seeded(
    total: Arc<AtomicU64>,
    interrupt: Arc<AtomicBool>,
    warmup: Arc<AtomicBool>,
    active_worker_count: Arc<AtomicUsize>,
) -> JoinHandle<()> {
    thread::spawn(move || {
        let metadata = v3::Metadata {
            designation: "",
            timestamp: &Timestamp::from_seconds(0),
            seed: &TEST_SEED,
        };

        let mut counter = 0u64;
        let mut warmup_cache = true;

        active_worker_count.fetch_add(1, Ordering::Relaxed);
        while !interrupt.load(Ordering::Relaxed) {
            if warmup_cache {
                warmup_cache = warmup.load(Ordering::Relaxed);
            }

            v3::generate(&metadata);
            if !warmup_cache {
                counter += 1;
            }
        }
        total.fetch_add(counter, Ordering::Relaxed);
        active_worker_count.fetch_sub(1, Ordering::Relaxed);
    })
}

fn run_unseeded(
    total: Arc<AtomicU64>,
    interrupt: Arc<AtomicBool>,
    warmup: Arc<AtomicBool>,
    active_worker_count: Arc<AtomicUsize>,
) -> JoinHandle<()> {
    thread::spawn(move || {
        let mut rng = rand::rng();
        let mut seed = [0u8; 32];
        let mut counter = 0u64;
        let mut warmup_cache = true;

        active_worker_count.fetch_add(1, Ordering::Relaxed);
        while !interrupt.load(Ordering::Relaxed) {
            if warmup_cache {
                warmup_cache = warmup.load(Ordering::Relaxed);
            }
            rng.fill_bytes(&mut seed);
            let metadata = v3::Metadata {
                designation: "",
                timestamp: &Timestamp::new(),
                seed: &seed,
            };
            let _ = v3::generate(&metadata);
            if !warmup_cache {
                counter += 1;
            }
        }
        total.fetch_add(counter, Ordering::Relaxed);
        active_worker_count.fetch_sub(1, Ordering::Relaxed);
    })
}

fn print_with_sep(number: u128) -> String {
    let mut n = number;
    let mut chunks: Vec<u16> = vec![];

    while n > 0 {
        chunks.push((n % 1000).try_into().unwrap());
        n /= 1000;
    }

    chunks
        .iter()
        .rev()
        .enumerate()
        .map(|(n, c)| {
            if n == 0 {
                format!("{c}")
            } else {
                format!("{c:03}")
            }
        })
        .collect::<Vec<String>>()
        .join(",")
}