//! Translated from libnyx/test.c

mod benchmark;
mod benchmark_baseline;
mod benchmark_nyx_no_pt;
mod benchmark_nyx_pt;
mod nyx;

use crate::benchmark::{Benchmark, ToCsv};
use crate::benchmark_baseline::Baseline;
use crate::benchmark_nyx_no_pt::BenchmarkNyxNoPt;
use crate::benchmark_nyx_pt::BenchmarkNyx;
use clap::Parser;
use std::fmt::Debug;
use std::fs;
use std::path::{Path, PathBuf};
use std::time::{Duration, Instant};

#[derive(Debug, Parser)]
#[command(about = "Tool to execute a program with nyx and trace it with intel pt")]
struct Args {
    /// Path to the client binary to execute
    client: PathBuf,
    /// Directory where the benchmark files get stored
    #[arg(short = 'o', long = "output", default_value = "benchmarks")]
    output_dir: PathBuf,
    /// Whether to only execute a single nyx run, used for testing
    #[arg(short = 'd', long = "debug", default_value = "false")]
    debug: bool,
}

fn main() -> anyhow::Result<()> {
    let args = Args::parse();

    if args.debug {
        dbg!(BenchmarkNyx::new(&args.client)?.execute_once());
        return Ok(());
    }

    let binary_name = args.client.file_name().unwrap().to_str().unwrap();
    {
        let benchmark_nyx = BenchmarkNyx::new(&args.client)?;
        benchmark(binary_name, &args.output_dir, benchmark_nyx)?;
    }
    {
        let benchmark_nyx_no_pt = BenchmarkNyxNoPt::new(&args.client)?;
        benchmark(binary_name, &args.output_dir, benchmark_nyx_no_pt)?;
    }
    benchmark(binary_name, &args.output_dir, Baseline(&args.client))?;

    Ok(())
}

fn benchmark<B: Benchmark>(
    binary_name: &str,
    output_dir: &Path,
    mut benchmark: B,
) -> anyhow::Result<()> {
    warmup(&mut benchmark);
    let results = benchmark_loop(&mut benchmark);
    write_results::<B>(binary_name, output_dir, &results)?;

    Ok(())
}

/// Warm up to make sure caches are initialized, etc.
fn warmup(benchmark: &mut impl Benchmark) {
    let warmup_duration = Duration::from_secs(5);
    println!("Warming up for {warmup_duration:?}");
    let mut iterations = 0;
    let start = Instant::now();
    while start.elapsed() < warmup_duration {
        benchmark.execute_once();
        iterations += 1;
    }
    println!("Warmed up for {iterations} iterations");
}

fn benchmark_loop<B: Benchmark>(benchmark: &mut B) -> Vec<B::BenchmarkResult> {
    let n = 50;
    println!("Perform {n} iterations...");
    (0..n)
        .map(|_| std::hint::black_box(benchmark.execute_once()))
        .collect::<Vec<_>>()
}

fn write_results<B: Benchmark>(
    binary_name: &str,
    output_dir: &Path,
    results: &[B::BenchmarkResult],
) -> anyhow::Result<()> {
    fs::create_dir_all(output_dir)?;
    let filename = format!("benchmark_{binary_name}_{}", B::TITLE);
    let file = fs::File::create(output_dir.join(filename))?;
    let mut writer = csv::WriterBuilder::new().from_writer(file);
    writer.write_record(B::BenchmarkResult::HEADERS)?;
    for result in results {
        result.write(&mut writer)?;
    }

    Ok(())
}