use rpki::repository::cert::Cert;
use rpki::repository::crl::Crl;
use rpki::repository::manifest::Manifest;
use rpki::repository::roa::Roa;
use rpki::repository::aspa::Aspa;
use rpki::repository::resources::{AsResources, IpResources};

use std::hint::black_box;
use std::path::{Path, PathBuf};
use std::time::Instant;

#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
enum ObjType {
    Cer,
    Crl,
    Manifest,
    Roa,
    Aspa,
}

impl ObjType {
    fn parse(s: &str) -> Result<Self, String> {
        match s {
            "cer" => Ok(Self::Cer),
            "crl" => Ok(Self::Crl),
            "manifest" => Ok(Self::Manifest),
            "roa" => Ok(Self::Roa),
            "aspa" => Ok(Self::Aspa),
            _ => Err("type must be one of: cer, crl, manifest, roa, aspa".into()),
        }
    }

    fn as_str(self) -> &'static str {
        match self {
            ObjType::Cer => "cer",
            ObjType::Crl => "crl",
            ObjType::Manifest => "manifest",
            ObjType::Roa => "roa",
            ObjType::Aspa => "aspa",
        }
    }

    fn ext(self) -> &'static str {
        match self {
            ObjType::Cer => "cer",
            ObjType::Crl => "crl",
            ObjType::Manifest => "mft",
            ObjType::Roa => "roa",
            ObjType::Aspa => "asa",
        }
    }
}

#[derive(Clone, Debug)]
struct Sample {
    obj_type: ObjType,
    name: String,
    path: PathBuf,
}

#[derive(Clone, Debug)]
struct Config {
    dir: PathBuf,
    type_filter: Option<ObjType>,
    sample_filter: Option<String>,
    fixed_iters: Option<u64>,
    warmup_iters: u64,
    rounds: u64,
    min_round_ms: u64,
    max_adaptive_iters: u64,
    strict: bool,
    cert_inspect: bool,
    out_csv: Option<PathBuf>,
    out_md: Option<PathBuf>,
}

fn usage_and_exit(err: Option<&str>) -> ! {
    if let Some(err) = err {
        eprintln!("error: {err}");
        eprintln!();
    }
    eprintln!(
        "Usage:\n\
  cargo run --release --manifest-path rpki/benchmark/routinator_object_bench/Cargo.toml -- [OPTIONS]\n\
\n\
Options:\n\
  --dir <PATH>            Fixtures root dir (default: ../../tests/benchmark/selected_der_v2)\n\
  --type <cer|crl|manifest|roa|aspa>    Filter by type\n\
  --sample <NAME>         Filter by sample name (e.g. p50)\n\
  --iters <N>             Fixed iterations per round (optional; otherwise adaptive)\n\
  --warmup-iters <N>      Warmup iterations (default: 50)\n\
  --rounds <N>            Rounds (default: 5)\n\
  --min-round-ms <MS>     Adaptive: minimum round time (default: 200)\n\
  --max-iters <N>         Adaptive: maximum iters (default: 1_000_000)\n\
  --strict <true|false>   Strict DER where applicable (default: true)\n\
  --cert-inspect          Also run Cert::inspect_ca/inspect_ee where applicable (default: false)\n\
  --out-csv <PATH>        Write CSV output\n\
  --out-md <PATH>         Write Markdown output\n\
"
    );
    std::process::exit(2);
}

fn parse_bool(s: &str, name: &str) -> bool {
    match s {
        "1" | "true" | "TRUE" | "yes" | "YES" => true,
        "0" | "false" | "FALSE" | "no" | "NO" => false,
        _ => usage_and_exit(Some(&format!("{name} must be true/false"))),
    }
}

fn parse_u64(s: &str, name: &str) -> u64 {
    s.parse::<u64>()
        .unwrap_or_else(|_| usage_and_exit(Some(&format!("{name} must be an integer"))))
}

fn default_samples_dir() -> PathBuf {
    PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../../tests/benchmark/selected_der_v2")
}

fn parse_args() -> Config {
    let mut dir: PathBuf = default_samples_dir();
    let mut type_filter: Option<ObjType> = None;
    let mut sample_filter: Option<String> = None;
    let mut fixed_iters: Option<u64> = None;
    let mut warmup_iters: u64 = 50;
    let mut rounds: u64 = 5;
    let mut min_round_ms: u64 = 200;
    let mut max_adaptive_iters: u64 = 1_000_000;
    let mut strict: bool = true;
    let mut cert_inspect: bool = false;
    let mut out_csv: Option<PathBuf> = None;
    let mut out_md: Option<PathBuf> = None;

    let mut args = std::env::args().skip(1);
    while let Some(arg) = args.next() {
        match arg.as_str() {
            "--dir" => dir = PathBuf::from(args.next().unwrap_or_else(|| usage_and_exit(None))),
            "--type" => {
                type_filter = Some(ObjType::parse(
                    &args.next().unwrap_or_else(|| usage_and_exit(None)),
                )
                .unwrap_or_else(|e| usage_and_exit(Some(&e))))
            }
            "--sample" => {
                sample_filter = Some(args.next().unwrap_or_else(|| usage_and_exit(None)))
            }
            "--iters" => {
                fixed_iters = Some(parse_u64(
                    &args.next().unwrap_or_else(|| usage_and_exit(None)),
                    "--iters",
                ))
            }
            "--warmup-iters" => {
                warmup_iters = parse_u64(
                    &args.next().unwrap_or_else(|| usage_and_exit(None)),
                    "--warmup-iters",
                )
            }
            "--rounds" => {
                rounds = parse_u64(&args.next().unwrap_or_else(|| usage_and_exit(None)), "--rounds")
            }
            "--min-round-ms" => {
                min_round_ms = parse_u64(
                    &args.next().unwrap_or_else(|| usage_and_exit(None)),
                    "--min-round-ms",
                )
            }
            "--max-iters" => {
                max_adaptive_iters = parse_u64(
                    &args.next().unwrap_or_else(|| usage_and_exit(None)),
                    "--max-iters",
                )
            }
            "--strict" => {
                strict = parse_bool(
                    &args.next().unwrap_or_else(|| usage_and_exit(None)),
                    "--strict",
                )
            }
            "--cert-inspect" => cert_inspect = true,
            "--out-csv" => out_csv = Some(PathBuf::from(args.next().unwrap_or_else(|| usage_and_exit(None)))),
            "--out-md" => out_md = Some(PathBuf::from(args.next().unwrap_or_else(|| usage_and_exit(None)))),
            "-h" | "--help" => usage_and_exit(None),
            _ => usage_and_exit(Some(&format!("unknown argument: {arg}"))),
        }
    }

    if warmup_iters == 0 {
        usage_and_exit(Some("--warmup-iters must be > 0"));
    }
    if rounds == 0 {
        usage_and_exit(Some("--rounds must be > 0"));
    }
    if min_round_ms == 0 {
        usage_and_exit(Some("--min-round-ms must be > 0"));
    }
    if max_adaptive_iters == 0 {
        usage_and_exit(Some("--max-iters must be > 0"));
    }
    if let Some(n) = fixed_iters {
        if n == 0 {
            usage_and_exit(Some("--iters must be > 0"));
        }
    }

    Config {
        dir,
        type_filter,
        sample_filter,
        fixed_iters,
        warmup_iters,
        rounds,
        min_round_ms,
        max_adaptive_iters,
        strict,
        cert_inspect,
        out_csv,
        out_md,
    }
}

fn read_samples(root: &Path) -> Vec<Sample> {
    let mut out = Vec::new();
    for obj_type in [
        ObjType::Cer,
        ObjType::Crl,
        ObjType::Manifest,
        ObjType::Roa,
        ObjType::Aspa,
    ] {
        let dir = root.join(obj_type.as_str());
        let rd = match std::fs::read_dir(&dir) {
            Ok(rd) => rd,
            Err(_) => continue,
        };
        for ent in rd.flatten() {
            let path = ent.path();
            if path.extension().and_then(|s| s.to_str()) != Some(obj_type.ext()) {
                continue;
            }
            let name = path
                .file_stem()
                .and_then(|s| s.to_str())
                .unwrap_or("unknown")
                .to_string();
            out.push(Sample { obj_type, name, path });
        }
    }
    out.sort_by(|a, b| a.obj_type.cmp(&b.obj_type).then_with(|| a.name.cmp(&b.name)));
    out
}

fn choose_iters_adaptive<F: FnMut()>(mut op: F, min_round_ms: u64, max_iters: u64) -> u64 {
    let min_secs = (min_round_ms as f64) / 1e3;
    let mut iters: u64 = 1;
    loop {
        let start = Instant::now();
        for _ in 0..iters {
            op();
        }
        let elapsed = start.elapsed().as_secs_f64();
        if elapsed >= min_secs {
            return iters;
        }
        if iters >= max_iters {
            return iters;
        }
        iters = (iters.saturating_mul(2)).min(max_iters);
    }
}

fn count_ip(res: &IpResources) -> u64 {
    if res.is_inherited() {
        return 1;
    }
    let Ok(blocks) = res.to_blocks() else {
        return 0;
    };
    blocks.iter().count() as u64
}

fn count_as(res: &AsResources) -> u64 {
    if res.is_inherited() {
        return 1;
    }
    let Ok(blocks) = res.to_blocks() else {
        return 0;
    };
    blocks.iter().count() as u64
}

fn complexity(obj_type: ObjType, bytes: &[u8], strict: bool, cert_inspect: bool) -> u64 {
    match obj_type {
        ObjType::Cer => {
            let cert = Cert::decode(bytes).expect("decode cert");
            if cert_inspect {
                if cert.is_ca() {
                    cert.inspect_ca(strict).expect("inspect ca");
                } else {
                    cert.inspect_ee(strict).expect("inspect ee");
                }
            }
            count_ip(cert.v4_resources())
                .saturating_add(count_ip(cert.v6_resources()))
                .saturating_add(count_as(cert.as_resources()))
        }
        ObjType::Crl => {
            let crl = Crl::decode(bytes).expect("decode crl");
            crl.revoked_certs().iter().count() as u64
        }
        ObjType::Manifest => {
            let mft = Manifest::decode(bytes, strict).expect("decode manifest");
            if cert_inspect {
                mft.cert().inspect_ee(strict).expect("inspect ee");
            }
            mft.content().len() as u64
        }
        ObjType::Roa => {
            let roa = Roa::decode(bytes, strict).expect("decode roa");
            if cert_inspect {
                roa.cert().inspect_ee(strict).expect("inspect ee");
            }
            roa.content().iter().count() as u64
        }
        ObjType::Aspa => {
            let asa = Aspa::decode(bytes, strict).expect("decode aspa");
            if cert_inspect {
                asa.cert().inspect_ee(strict).expect("inspect ee");
            }
            asa.content().provider_as_set().len() as u64
        }
    }
}

fn decode_profile(obj_type: ObjType, bytes: &[u8], strict: bool, cert_inspect: bool) {
    match obj_type {
        ObjType::Cer => {
            let cert = Cert::decode(black_box(bytes)).expect("decode cert");
            if cert_inspect {
                if cert.is_ca() {
                    cert.inspect_ca(strict).expect("inspect ca");
                } else {
                    cert.inspect_ee(strict).expect("inspect ee");
                }
            }
            black_box(cert);
        }
        ObjType::Crl => {
            let crl = Crl::decode(black_box(bytes)).expect("decode crl");
            black_box(crl);
        }
        ObjType::Manifest => {
            let mft = Manifest::decode(black_box(bytes), strict).expect("decode manifest");
            if cert_inspect {
                mft.cert().inspect_ee(strict).expect("inspect ee");
            }
            black_box(mft);
        }
        ObjType::Roa => {
            let roa = Roa::decode(black_box(bytes), strict).expect("decode roa");
            if cert_inspect {
                roa.cert().inspect_ee(strict).expect("inspect ee");
            }
            black_box(roa);
        }
        ObjType::Aspa => {
            let asa = Aspa::decode(black_box(bytes), strict).expect("decode aspa");
            if cert_inspect {
                asa.cert().inspect_ee(strict).expect("inspect ee");
            }
            black_box(asa);
        }
    }
}

#[derive(Clone, Debug)]
struct ResultRow {
    obj_type: String,
    sample: String,
    size_bytes: usize,
    complexity: u64,
    avg_ns_per_op: f64,
    ops_per_sec: f64,
}

fn render_markdown(title: &str, rows: &[ResultRow]) -> String {
    let mut out = String::new();
    out.push_str(&format!("# {title}\n\n"));
    out.push_str("| type | sample | size_bytes | complexity | avg ns/op | ops/s |\n");
    out.push_str("|---|---|---:|---:|---:|---:|\n");
    for r in rows {
        out.push_str(&format!(
            "| {} | {} | {} | {} | {:.2} | {:.2} |\n",
            r.obj_type, r.sample, r.size_bytes, r.complexity, r.avg_ns_per_op, r.ops_per_sec
        ));
    }
    out
}

fn render_csv(rows: &[ResultRow]) -> String {
    let mut out = String::new();
    out.push_str("type,sample,size_bytes,complexity,avg_ns_per_op,ops_per_sec\n");
    for r in rows {
        let sample = r.sample.replace('"', "\"\"");
        out.push_str(&format!(
            "{},{},{},{},{:.6},{:.6}\n",
            r.obj_type,
            format!("\"{}\"", sample),
            r.size_bytes,
            r.complexity,
            r.avg_ns_per_op,
            r.ops_per_sec
        ));
    }
    out
}

fn create_parent_dirs(path: &Path) {
    if let Some(parent) = path.parent() {
        std::fs::create_dir_all(parent).unwrap_or_else(|e| {
            panic!("create_dir_all {}: {e}", parent.display());
        });
    }
}

fn write_text_file(path: &Path, content: &str) {
    create_parent_dirs(path);
    std::fs::write(path, content).unwrap_or_else(|e| panic!("write {}: {e}", path.display()));
}

fn main() {
    let cfg = parse_args();
    let mut samples = read_samples(&cfg.dir);
    if samples.is_empty() {
        usage_and_exit(Some(&format!(
            "no samples found under: {}",
            cfg.dir.display()
        )));
    }

    if let Some(t) = cfg.type_filter {
        samples.retain(|s| s.obj_type == t);
        if samples.is_empty() {
            usage_and_exit(Some(&format!("no sample matched --type {}", t.as_str())));
        }
    }
    if let Some(filter) = cfg.sample_filter.as_deref() {
        samples.retain(|s| s.name == filter);
        if samples.is_empty() {
            usage_and_exit(Some(&format!("no sample matched --sample {filter}")));
        }
    }

    println!("# Routinator baseline (rpki crate) decode benchmark (selected_der_v2)");
    println!();
    println!("- dir: {}", cfg.dir.display());
    println!("- strict: {}", cfg.strict);
    println!("- cert_inspect: {}", cfg.cert_inspect);
    if let Some(t) = cfg.type_filter {
        println!("- type: {}", t.as_str());
    }
    if let Some(s) = cfg.sample_filter.as_deref() {
        println!("- sample: {}", s);
    }
    if let Some(n) = cfg.fixed_iters {
        println!("- iters: {} (fixed)", n);
    } else {
        println!(
            "- warmup: {} iters, rounds: {}, min_round: {}ms (adaptive iters, max {})",
            cfg.warmup_iters, cfg.rounds, cfg.min_round_ms, cfg.max_adaptive_iters
        );
    }
    if let Some(p) = cfg.out_csv.as_ref() {
        println!("- out_csv: {}", p.display());
    }
    if let Some(p) = cfg.out_md.as_ref() {
        println!("- out_md: {}", p.display());
    }
    println!();

    println!("| type | sample | size_bytes | complexity | avg ns/op | ops/s |");
    println!("|---|---|---:|---:|---:|---:|");

    let mut rows: Vec<ResultRow> = Vec::with_capacity(samples.len());
    for sample in &samples {
        let bytes = std::fs::read(&sample.path)
            .unwrap_or_else(|e| panic!("read {}: {e}", sample.path.display()));
        let size_bytes = bytes.len();
        let complexity = complexity(sample.obj_type, bytes.as_slice(), cfg.strict, cfg.cert_inspect);

        for _ in 0..cfg.warmup_iters {
            decode_profile(sample.obj_type, bytes.as_slice(), cfg.strict, cfg.cert_inspect);
        }

        let mut per_round_ns_per_op = Vec::with_capacity(cfg.rounds as usize);
        for _round in 0..cfg.rounds {
            let iters = if let Some(n) = cfg.fixed_iters {
                n
            } else {
                choose_iters_adaptive(
                    || decode_profile(sample.obj_type, bytes.as_slice(), cfg.strict, cfg.cert_inspect),
                    cfg.min_round_ms,
                    cfg.max_adaptive_iters,
                )
            };
            let start = Instant::now();
            for _ in 0..iters {
                decode_profile(sample.obj_type, bytes.as_slice(), cfg.strict, cfg.cert_inspect);
            }
            let elapsed = start.elapsed();
            let total_ns = elapsed.as_secs_f64() * 1e9;
            per_round_ns_per_op.push(total_ns / (iters as f64));
        }

        let avg_ns = per_round_ns_per_op.iter().sum::<f64>() / (per_round_ns_per_op.len() as f64);
        let ops_per_sec = 1e9_f64 / avg_ns;

        println!(
            "| {} | {} | {} | {} | {:.2} | {:.2} |",
            sample.obj_type.as_str(),
            sample.name,
            size_bytes,
            complexity,
            avg_ns,
            ops_per_sec
        );

        rows.push(ResultRow {
            obj_type: sample.obj_type.as_str().to_string(),
            sample: sample.name.clone(),
            size_bytes,
            complexity,
            avg_ns_per_op: avg_ns,
            ops_per_sec,
        });
    }

    if let Some(path) = cfg.out_md.as_ref() {
        let md = render_markdown(
            "Routinator baseline (rpki crate) decode+inspect (selected_der_v2)",
            &rows,
        );
        write_text_file(path, &md);
        eprintln!("Wrote {}", path.display());
    }
    if let Some(path) = cfg.out_csv.as_ref() {
        let csv = render_csv(&rows);
        write_text_file(path, &csv);
        eprintln!("Wrote {}", path.display());
    }
}