i0 0.1.4

use i0::{Timestamp, v3};
use rand::RngCore;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicU8, AtomicU64, Ordering};
use std::thread::{self, JoinHandle};
use std::time::{Duration, Instant};

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,
];

fn main() {
    let nproc = thread::available_parallelism().unwrap().get();
    let test_cases = [
        // seeded?, #threads, warmup (ms), duration (ms)
        (false, 1, 1000, 3000),
        (false, nproc, 2000, 5000),
        (true, nproc, 1000, 3000),
    ];

    for (seeded, threads, t_warm, t_run) in test_cases {
        run(
            seeded,
            threads,
            Duration::from_millis(t_warm),
            Duration::from_millis(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(AtomicU8::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:>2}]",
            if seeded { "seeded" } else { "unseeded" }
        );
        // warmup
        println!("{prefix} warming up ({t_warm:?})");
        thread::sleep(t_warm);
        // start
        print!("\x1b[F");
        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} {:>7} 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<AtomicU8>,
) -> 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<AtomicU8>,
) -> 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(",")
}