use rpki::data_model::manifest::ManifestObject;
use std::hint::black_box;
use std::path::PathBuf;
use std::time::Instant;

#[derive(Debug, Clone)]
struct Config {
    sample: Option<String>,
    manifest_path: Option<PathBuf>,
    iterations: u64,
    warmup_iterations: u64,
    repeats: u32,
}

fn usage_and_exit() -> ! {
    eprintln!(
        "Usage:\n  ours-manifest-bench (--sample <name> | --manifest <path>) [--iterations N] [--warmup-iterations N] [--repeats N]\n\nExamples:\n  cargo run --release -- --sample small-01 --iterations 20000 --warmup-iterations 2000 --repeats 3\n  cargo run --release -- --manifest ../../tests/benchmark/selected_der/small-01.mft"
    );
    std::process::exit(2);
}

fn parse_args() -> Config {
    let mut sample: Option<String> = None;
    let mut manifest_path: Option<PathBuf> = None;
    let mut iterations: u64 = 20_000;
    let mut warmup_iterations: u64 = 2_000;
    let mut repeats: u32 = 3;

    let mut args = std::env::args().skip(1);
    while let Some(arg) = args.next() {
        match arg.as_str() {
            "--sample" => sample = Some(args.next().unwrap_or_else(|| usage_and_exit())),
            "--manifest" => {
                manifest_path = Some(PathBuf::from(args.next().unwrap_or_else(|| usage_and_exit())))
            }
            "--iterations" => {
                iterations = args
                    .next()
                    .unwrap_or_else(|| usage_and_exit())
                    .parse()
                    .unwrap_or_else(|_| usage_and_exit())
            }
            "--warmup-iterations" => {
                warmup_iterations = args
                    .next()
                    .unwrap_or_else(|| usage_and_exit())
                    .parse()
                    .unwrap_or_else(|_| usage_and_exit())
            }
            "--repeats" => {
                repeats = args
                    .next()
                    .unwrap_or_else(|| usage_and_exit())
                    .parse()
                    .unwrap_or_else(|_| usage_and_exit())
            }
            "-h" | "--help" => usage_and_exit(),
            _ => usage_and_exit(),
        }
    }

    if sample.is_none() && manifest_path.is_none() {
        usage_and_exit();
    }
    if sample.is_some() && manifest_path.is_some() {
        usage_and_exit();
    }

    Config {
        sample,
        manifest_path,
        iterations,
        warmup_iterations,
        repeats,
    }
}

fn derive_manifest_path(sample: &str) -> PathBuf {
    // Assumes current working directory is `rpki/benchmark/ours_manifest_bench`.
    PathBuf::from(format!("../../tests/benchmark/selected_der/{sample}.mft"))
}

fn main() {
    let cfg = parse_args();
    let manifest_path = cfg
        .manifest_path
        .clone()
        .unwrap_or_else(|| derive_manifest_path(cfg.sample.as_deref().unwrap()));

    let bytes = std::fs::read(&manifest_path).unwrap_or_else(|e| {
        eprintln!("read manifest fixture failed: {e}; path={}", manifest_path.display());
        std::process::exit(1);
    });

    let decoded_once = ManifestObject::decode_der(&bytes).unwrap_or_else(|e| {
        eprintln!("decode failed: {e}; path={}", manifest_path.display());
        std::process::exit(1);
    });
    let file_count = decoded_once.manifest.file_count();

    let mut round_ns_per_op: Vec<f64> = Vec::with_capacity(cfg.repeats as usize);
    let mut round_ops_per_s: Vec<f64> = Vec::with_capacity(cfg.repeats as usize);

    for _round in 0..cfg.repeats {
        for _ in 0..cfg.warmup_iterations {
            let obj = ManifestObject::decode_der(black_box(&bytes)).expect("warmup decode");
            black_box(obj);
        }

        let start = Instant::now();
        for _ in 0..cfg.iterations {
            let obj = ManifestObject::decode_der(black_box(&bytes)).expect("timed decode");
            black_box(obj);
        }
        let elapsed = start.elapsed();

        let ns_per_op = (elapsed.as_secs_f64() * 1e9) / (cfg.iterations as f64);
        let ops_per_s = (cfg.iterations as f64) / elapsed.as_secs_f64();
        round_ns_per_op.push(ns_per_op);
        round_ops_per_s.push(ops_per_s);
    }

    let avg_ns_per_op = round_ns_per_op.iter().sum::<f64>() / (round_ns_per_op.len() as f64);
    let avg_ops_per_s = round_ops_per_s.iter().sum::<f64>() / (round_ops_per_s.len() as f64);

    let sample_name = cfg.sample.clone().unwrap_or_else(|| {
        manifest_path
            .file_name()
            .map(|s| s.to_string_lossy().to_string())
            .unwrap_or_else(|| manifest_path.display().to_string())
    });
    let sample_name = sample_name
        .strip_suffix(".mft")
        .unwrap_or(&sample_name)
        .to_string();

    println!("fixture: {}", manifest_path.display());
    println!();
    println!("| sample | avg ns/op | ops/s | file count |");
    println!("|---|---:|---:|---:|");
    println!(
        "| {} | {:.2} | {:.2} | {} |",
        sample_name, avg_ns_per_op, avg_ops_per_s, file_count
    );
}