i0 0.2.2

use clap::Parser;
use i0::{Timestamp, v3};
use rand::Rng;
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");
const BIN_NAME: &str = env!("CARGO_BIN_NAME");
const TARGET_OS: &str = cfg_select! {
    windows => "windows",
    target_os = "macos" => "macos",
    unix => "unix",
    _ => "unknown",
};
const TARGET_FEATURE: &str = cfg_select! {
  all(target_feature = "avx512bw", target_feature = "avx512cd", target_feature = "avx512dq", target_feature = "avx512vl") => "x86-64-v4",
  all(target_feature = "avx2", target_feature = "bmi1", target_feature = "bmi2") => "x86-64-v3",
  all(target_feature = "sse4.2", target_feature = "ssse3") => "x64-64-v2",
  all(target_feature = "sse2") => "x64-64-v1",
  all(target_feature = "sve2", target_feature = "bf16") => "armv9.1",
  all(target_feature = "dotprod", target_feature = "fcma", target_feature = "fhm", target_feature = "flagm", target_feature = "fp16", target_feature = "frintts", target_feature = "jsconv", target_feature = "rdm", target_feature = "sve2") => "armv9",
  all(target_feature = "bf16", target_feature = "i8mm") => "armv8.6",
  all(target_feature = "bti", target_feature = "dpb2", target_feature = "frintts", target_feature = "sb", target_feature = "sve") => "armv8.5",
  all(target_feature = "dit", target_feature = "dotprod", target_feature = "flagm", target_feature = "rcpc2") => "armv8.4",
  any(target_feature = "bti", target_feature = "frintts", target_feature = "mte", target_feature = "rand") => "armv8.4",
  all(target_feature = "fcma", target_feature = "jsconv", target_feature = "rcpc", any(target_feature = "paca", target_feature = "pacg")) => "armv8.3",
  any(target_feature = "dit") => "armv8.3",
  all(target_feature = "dpb", target_feature = "pan", target_feature = "ras") => "armv8.2",
  any(target_feature = "bf16", target_feature = "f32mm", target_feature = "f64mm", target_feature = "fcma", target_feature = "fp16", target_feature = "jsconv", target_feature = "paca", target_feature = "pacg", target_feature = "rcpc", target_feature = "sve") => "armv8.2",
  all(target_feature = "crc", target_feature = "lor", target_feature = "lse", target_feature = "pan", target_feature = "rdm") => "armv8.1",
  any(target_feature = "dotprod", target_feature = "dpb", target_feature = "fhm", target_feature = "flagm", target_feature = "i8mm", target_feature = "sha3", target_feature = "sm4", target_feature = "spe") => "armv8.1",
  any(target_feature = "aes", target_feature = "crc", target_feature = "lor", target_feature = "lse", target_feature = "neon", target_feature = "pan", target_feature = "pmuv3", target_feature = "ras", target_feature = "rdm", target_feature = "sb", target_feature = "sha2", target_feature = "ssbs", target_feature = "vh") => "armv8",
  _ => "",
};

#[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 multi-threaded test before measuring throughput
    #[arg(short, long, default_value = "2s")]
    warmup_multi: humantime::Duration,

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

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

    /// how long to run the single-threaded test for (excluding warmup time)
    #[arg(short = 'T', long, default_value = "5s")]
    time_single: 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_multi_warmup = *cli.warmup_multi;
    let t_multi = *cli.time_multi;
    let t_single_warmup = *cli.warmup_single;
    let t_single = *cli.time_single;
    if t_multi.is_zero() || t_single.is_zero() {
        eprintln!("don't be silly, you can't run a benchmark with a 0s duration!");
        exit(1)
    }

    print_debug_info();

    if nproc != 1 {
        run(cli.seeded, 1, t_single_warmup, t_single);
    }

    run(cli.seeded, nproc, t_multi_warmup, t_multi);
}

fn print_debug_info() {
    println!("{BIN_NAME} 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"
        };
        println!("host: {TARGET_OS} (\"{host}\")");
        println!("cpu: {cpu} ({TARGET_FEATURE})");
    } else {
        println!("host: {TARGET_OS}");
        println!("cpu: unknown ({TARGET_FEATURE})")
    }
}

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} op/s ({t_run:?}) {}",
            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(",")
}