use rpki::data_model::aspa::AspaObject;
use rpki::data_model::crl::RpkixCrl;
use rpki::data_model::manifest::ManifestObject;
use rpki::data_model::rc::{
    AsIdentifierChoice, AsResourceSet, IpAddressChoice, IpResourceSet, ResourceCertificate,
};
use rpki::data_model::roa::RoaObject;

use rpki::storage::PackFile;
use rpki::storage::RawByHashEntry;
use rpki::storage::RocksStore;
use sha2::Digest;

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 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,
}

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

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 env_u64(name: &str, default: u64) -> u64 {
    std::env::var(name)
        .ok()
        .and_then(|s| s.parse::<u64>().ok())
        .unwrap_or(default)
}

fn env_u64_opt(name: &str) -> Option<u64> {
    std::env::var(name).ok().and_then(|s| s.parse::<u64>().ok())
}

fn env_bool(name: &str) -> bool {
    matches!(
        std::env::var(name).as_deref(),
        Ok("1") | Ok("true") | Ok("TRUE") | Ok("yes") | Ok("YES")
    )
}

fn env_string(name: &str) -> Option<String> {
    std::env::var(name).ok().filter(|s| !s.trim().is_empty())
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum BenchMode {
    Both,
    DecodeValidate,
    Landing,
}

impl BenchMode {
    fn parse(s: &str) -> Result<Self, String> {
        match s {
            "" | "both" => Ok(Self::Both),
            "decode" | "decode_validate" | "decode+validate" => Ok(Self::DecodeValidate),
            "landing" => Ok(Self::Landing),
            _ => Err(format!(
                "invalid BENCH_MODE='{s}', expected one of: both, decode_validate, landing"
            )),
        }
    }

    fn do_decode(self) -> bool {
        matches!(self, BenchMode::Both | BenchMode::DecodeValidate)
    }

    fn do_landing(self) -> bool {
        matches!(self, BenchMode::Both | BenchMode::Landing)
    }

    fn as_str(self) -> &'static str {
        match self {
            BenchMode::Both => "both",
            BenchMode::DecodeValidate => "decode_validate",
            BenchMode::Landing => "landing",
        }
    }
}

fn escape_md(s: &str) -> String {
    s.replace('|', "\\|").replace('\n', " ")
}

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 fmt_rfc3339_utc_now() -> String {
    use time::format_description::well_known::Rfc3339;
    time::OffsetDateTime::now_utc()
        .to_offset(time::UtcOffset::UTC)
        .format(&Rfc3339)
        .unwrap_or_else(|_| "unknown".to_string())
}

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 ip_resource_count(set: &IpResourceSet) -> u64 {
    let mut n = 0u64;
    for fam in &set.families {
        match &fam.choice {
            IpAddressChoice::Inherit => n = n.saturating_add(1),
            IpAddressChoice::AddressesOrRanges(items) => n = n.saturating_add(items.len() as u64),
        }
    }
    n
}

fn as_choice_count(choice: &AsIdentifierChoice) -> u64 {
    match choice {
        AsIdentifierChoice::Inherit => 1,
        AsIdentifierChoice::AsIdsOrRanges(items) => items.len() as u64,
    }
}

fn as_resource_count(set: &AsResourceSet) -> u64 {
    let mut n = 0u64;
    if let Some(c) = set.asnum.as_ref() {
        n = n.saturating_add(as_choice_count(c));
    }
    if let Some(c) = set.rdi.as_ref() {
        n = n.saturating_add(as_choice_count(c));
    }
    n
}

fn complexity_main(obj_type: ObjType, bytes: &[u8]) -> u64 {
    match obj_type {
        ObjType::Cer => {
            let cert = ResourceCertificate::decode_der(bytes).expect("decode cert");
            let ip = cert
                .tbs
                .extensions
                .ip_resources
                .as_ref()
                .map(ip_resource_count)
                .unwrap_or(0);
            let asn = cert
                .tbs
                .extensions
                .as_resources
                .as_ref()
                .map(as_resource_count)
                .unwrap_or(0);
            ip.saturating_add(asn)
        }
        ObjType::Crl => {
            let crl = RpkixCrl::decode_der(bytes).expect("decode crl");
            crl.revoked_certs.len() as u64
        }
        ObjType::Manifest => {
            let mft = ManifestObject::decode_der(bytes).expect("decode manifest");
            mft.manifest.file_count() as u64
        }
        ObjType::Roa => {
            let roa = RoaObject::decode_der(bytes).expect("decode roa");
            let mut n = 0u64;
            for fam in &roa.roa.ip_addr_blocks {
                n = n.saturating_add(fam.addresses.len() as u64);
            }
            n
        }
        ObjType::Aspa => {
            let asa = AspaObject::decode_der(bytes).expect("decode aspa");
            asa.aspa.provider_as_ids.len() as u64
        }
    }
}

fn decode_validate(obj_type: ObjType, bytes: &[u8]) {
    match obj_type {
        ObjType::Cer => {
            let decoded =
                ResourceCertificate::decode_der(std::hint::black_box(bytes)).expect("decode cert");
            std::hint::black_box(decoded);
        }
        ObjType::Crl => {
            let decoded = RpkixCrl::decode_der(std::hint::black_box(bytes)).expect("decode crl");
            std::hint::black_box(decoded);
        }
        ObjType::Manifest => {
            let decoded =
                ManifestObject::decode_der(std::hint::black_box(bytes)).expect("decode manifest");
            std::hint::black_box(decoded);
        }
        ObjType::Roa => {
            let decoded = RoaObject::decode_der(std::hint::black_box(bytes)).expect("decode roa");
            std::hint::black_box(decoded);
        }
        ObjType::Aspa => {
            let decoded = AspaObject::decode_der(std::hint::black_box(bytes)).expect("decode aspa");
            std::hint::black_box(decoded);
        }
    }
}

fn landing_packfile_cbor_put(store: &RocksStore, obj_type: ObjType, sample: &str, bytes: &[u8]) {
    #[derive(serde::Serialize)]
    struct BenchPackFile<'a> {
        rsync_uri: &'a str,
        bytes: &'a [u8],
        sha256: [u8; 32],
    }

    let rsync_uri = format!(
        "rsync://bench.invalid/{}/{}.{}",
        obj_type.as_str(),
        sample,
        obj_type.ext()
    );
    let pf = PackFile::from_bytes_compute_sha256(rsync_uri, bytes.to_vec());
    let bench_pf = BenchPackFile {
        rsync_uri: &pf.rsync_uri,
        bytes: pf.bytes().expect("load packfile bytes"),
        sha256: pf.sha256,
    };
    let encoded =
        serde_cbor::to_vec(std::hint::black_box(&bench_pf)).expect("cbor encode packfile");
    let key = format!("bench:packfile:{}:{}", obj_type.as_str(), sample);
    let sha256_hex = hex::encode(sha2::Sha256::digest(&encoded));
    let mut entry = RawByHashEntry::from_bytes(sha256_hex, encoded);
    entry.origin_uris.push(key);
    entry.object_type = Some("cbor".to_string());
    entry.encoding = Some("cbor".to_string());
    store
        .put_raw_by_hash_entry(&entry)
        .expect("store raw_by_hash");
}

#[derive(Clone, Debug, serde::Serialize)]
struct RunConfig {
    dir: String,
    mode: String,
    obj_type: Option<String>,
    sample: Option<String>,
    fixed_iters: Option<u64>,
    warmup_iters: u64,
    rounds: u64,
    min_round_ms: u64,
    max_adaptive_iters: u64,
    timestamp_utc: String,
}

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

#[derive(Clone, Debug, serde::Serialize)]
struct BenchmarkOutput {
    config: RunConfig,
    rows: Vec<ResultRow>,
}

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",
            escape_md(&r.obj_type),
            escape_md(&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 {
        // sample names are controlled by our fixtures; keep it simple but safe for commas/quotes.
        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
}

#[test]
#[ignore = "manual performance benchmark; decode+validate and landing for selected_der_v2"]
fn stage2_decode_validate_and_landing_benchmark_selected_der_v2() {
    let dir = env_string("BENCH_DIR")
        .map(PathBuf::from)
        .unwrap_or_else(default_samples_dir);

    let mode = env_string("BENCH_MODE")
        .as_deref()
        .map(BenchMode::parse)
        .transpose()
        .unwrap_or_else(|e| panic!("{e}"))
        .unwrap_or(BenchMode::Both);

    let type_filter = env_string("BENCH_TYPE");
    let sample_filter = env_string("BENCH_SAMPLE");
    let fixed_iters = env_u64_opt("BENCH_ITERS");
    let warmup_iters = env_u64("BENCH_WARMUP_ITERS", 50);
    let rounds = env_u64("BENCH_ROUNDS", 5);
    let min_round_ms = env_u64("BENCH_MIN_ROUND_MS", 200);
    let max_adaptive_iters = env_u64("BENCH_MAX_ITERS", 1_000_000);
    let verbose = env_bool("BENCH_VERBOSE");
    let out_md = env_string("BENCH_OUT_MD").map(PathBuf::from);
    let out_json = env_string("BENCH_OUT_JSON").map(PathBuf::from);
    let out_csv = env_string("BENCH_OUT_CSV").map(PathBuf::from);
    let out_md_landing = env_string("BENCH_OUT_MD_LANDING").map(PathBuf::from);
    let out_json_landing = env_string("BENCH_OUT_JSON_LANDING").map(PathBuf::from);
    let out_csv_landing = env_string("BENCH_OUT_CSV_LANDING").map(PathBuf::from);

    if let Some(n) = fixed_iters {
        assert!(n >= 1, "BENCH_ITERS must be >= 1");
    }
    assert!(rounds >= 1, "BENCH_ROUNDS must be >= 1");
    assert!(min_round_ms >= 1, "BENCH_MIN_ROUND_MS must be >= 1");
    assert!(max_adaptive_iters >= 1, "BENCH_MAX_ITERS must be >= 1");

    let mut samples = read_samples(&dir);
    assert!(
        !samples.is_empty(),
        "no samples found under: {}",
        dir.display()
    );

    if let Some(t) = type_filter.as_deref() {
        samples.retain(|s| s.obj_type.as_str() == t);
        assert!(!samples.is_empty(), "no sample matched BENCH_TYPE={t}");
    }
    if let Some(filter) = sample_filter.as_deref() {
        samples.retain(|s| s.name == filter);
        assert!(
            !samples.is_empty(),
            "no sample matched BENCH_SAMPLE={filter}"
        );
    }

    println!("# Stage2 decode+validate benchmark (selected_der_v2)");
    println!();
    println!("- dir: {}", dir.display());
    println!("- mode: {}", mode.as_str());
    if let Some(t) = type_filter.as_deref() {
        println!("- type: {}", t);
    }
    if let Some(s) = sample_filter.as_deref() {
        println!("- sample: {}", s);
    }
    if let Some(n) = fixed_iters {
        println!("- iters: {} (fixed)", n);
    } else {
        println!(
            "- warmup: {} iters, rounds: {}, min_round: {}ms (adaptive iters, max {})",
            warmup_iters, rounds, min_round_ms, max_adaptive_iters
        );
    }
    if verbose {
        println!("- verbose: true");
    }
    if let Some(p) = out_md.as_ref() {
        println!("- out_md: {}", p.display());
    }
    if let Some(p) = out_json.as_ref() {
        println!("- out_json: {}", p.display());
    }
    if let Some(p) = out_csv.as_ref() {
        println!("- out_csv: {}", p.display());
    }
    if let Some(p) = out_md_landing.as_ref() {
        println!("- out_md_landing: {}", p.display());
    }
    if let Some(p) = out_json_landing.as_ref() {
        println!("- out_json_landing: {}", p.display());
    }
    if let Some(p) = out_csv_landing.as_ref() {
        println!("- out_csv_landing: {}", p.display());
    }
    println!();

    if mode.do_decode() {
        println!("## decode+validate");
        println!();
        println!("| type | sample | size_bytes | complexity | avg ns/op | ops/s |");
        println!("|---|---|---:|---:|---:|---:|");
    }

    let mut decode_rows: Vec<ResultRow> = Vec::with_capacity(samples.len());

    let temp = tempfile::tempdir().expect("tempdir");
    let store = RocksStore::open(temp.path()).expect("open rocksdb");

    let mut landing_rows: Vec<ResultRow> = Vec::with_capacity(samples.len());

    for s in &samples {
        let bytes =
            std::fs::read(&s.path).unwrap_or_else(|e| panic!("read {}: {e}", s.path.display()));
        let size_bytes = bytes.len();
        let complexity = if mode.do_decode() {
            complexity_main(s.obj_type, bytes.as_slice())
        } else {
            0
        };

        if mode.do_decode() {
            // Warm-up decode+validate.
            for _ in 0..warmup_iters {
                decode_validate(s.obj_type, bytes.as_slice());
            }

            let mut per_round_ns_per_op = Vec::with_capacity(rounds as usize);
            for round in 0..rounds {
                let iters = if let Some(n) = fixed_iters {
                    n
                } else {
                    choose_iters_adaptive(
                        || decode_validate(s.obj_type, bytes.as_slice()),
                        min_round_ms,
                        max_adaptive_iters,
                    )
                };

                let start = Instant::now();
                for _ in 0..iters {
                    decode_validate(s.obj_type, bytes.as_slice());
                }
                let elapsed = start.elapsed();
                let total_ns = elapsed.as_secs_f64() * 1e9;
                let ns_per_op = total_ns / (iters as f64);
                per_round_ns_per_op.push(ns_per_op);

                if verbose {
                    println!(
                        "# {}.{} round {}: iters={} total_ms={:.2} ns/op={:.2}",
                        s.obj_type.as_str(),
                        s.name,
                        round + 1,
                        iters,
                        elapsed.as_secs_f64() * 1e3,
                        ns_per_op
                    );
                }
            }

            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} |",
                s.obj_type.as_str(),
                s.name,
                size_bytes,
                complexity,
                avg_ns,
                ops_per_sec
            );

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

        if mode.do_landing() {
            // Landing benchmark: PackFile(from_bytes_compute_sha256) + CBOR + RocksDB current-state landing.
            let mut per_round_ns_per_op = Vec::with_capacity(rounds as usize);
            for _ in 0..warmup_iters {
                landing_packfile_cbor_put(&store, s.obj_type, &s.name, bytes.as_slice());
            }
            for round in 0..rounds {
                let iters = if let Some(n) = fixed_iters {
                    n
                } else {
                    choose_iters_adaptive(
                        || landing_packfile_cbor_put(&store, s.obj_type, &s.name, bytes.as_slice()),
                        min_round_ms,
                        max_adaptive_iters,
                    )
                };

                let start = Instant::now();
                for _ in 0..iters {
                    landing_packfile_cbor_put(&store, s.obj_type, &s.name, bytes.as_slice());
                }
                let elapsed = start.elapsed();
                let total_ns = elapsed.as_secs_f64() * 1e9;
                let ns_per_op = total_ns / (iters as f64);
                per_round_ns_per_op.push(ns_per_op);

                if verbose {
                    println!(
                        "# landing {}.{} round {}: iters={} total_ms={:.2} ns/op={:.2}",
                        s.obj_type.as_str(),
                        s.name,
                        round + 1,
                        iters,
                        elapsed.as_secs_f64() * 1e3,
                        ns_per_op
                    );
                }
            }

            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;

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

    if mode.do_decode() {
        println!();
    }
    if mode.do_landing() {
        println!(
            "## landing (PackFile::from_bytes_compute_sha256 + CBOR + RocksDB current-state landing)"
        );
        println!();
        println!("| type | sample | size_bytes | complexity | avg ns/op | ops/s |");
        println!("|---|---|---:|---:|---:|---:|");
        for r in &landing_rows {
            println!(
                "| {} | {} | {} | {} | {:.2} | {:.2} |",
                escape_md(&r.obj_type),
                escape_md(&r.sample),
                r.size_bytes,
                r.complexity,
                r.avg_ns_per_op,
                r.ops_per_sec
            );
        }
        println!();
    }

    if mode.do_decode() && (out_md.is_some() || out_json.is_some() || out_csv.is_some()) {
        let timestamp_utc = fmt_rfc3339_utc_now();
        let cfg = RunConfig {
            dir: dir.display().to_string(),
            mode: "decode_validate".to_string(),
            obj_type: type_filter.clone(),
            sample: sample_filter.clone(),
            fixed_iters,
            warmup_iters,
            rounds,
            min_round_ms,
            max_adaptive_iters,
            timestamp_utc,
        };

        if let Some(path) = out_md {
            let md = render_markdown(
                "Stage2 decode+validate benchmark (selected_der_v2)",
                &decode_rows,
            );
            write_text_file(&path, &md);
            eprintln!("Wrote {}", path.display());
        }
        if let Some(path) = out_json {
            let json = serde_json::to_string_pretty(&BenchmarkOutput {
                config: cfg,
                rows: decode_rows.clone(),
            })
            .expect("serialize json");
            write_text_file(&path, &json);
            eprintln!("Wrote {}", path.display());
        }
        if let Some(path) = out_csv {
            let csv = render_csv(&decode_rows);
            write_text_file(&path, &csv);
            eprintln!("Wrote {}", path.display());
        }
    }

    if mode.do_landing()
        && (out_md_landing.is_some() || out_json_landing.is_some() || out_csv_landing.is_some())
    {
        let timestamp_utc = fmt_rfc3339_utc_now();
        let cfg = RunConfig {
            dir: dir.display().to_string(),
            mode: "landing_packfile_cbor_put".to_string(),
            obj_type: type_filter,
            sample: sample_filter,
            fixed_iters,
            warmup_iters,
            rounds,
            min_round_ms,
            max_adaptive_iters,
            timestamp_utc,
        };

        if let Some(path) = out_md_landing {
            let md = render_markdown(
                "Stage2 landing benchmark (PackFile CBOR + RocksDB current-state landing)",
                &landing_rows,
            );
            write_text_file(&path, &md);
            eprintln!("Wrote {}", path.display());
        }
        if let Some(path) = out_json_landing {
            let json = serde_json::to_string_pretty(&BenchmarkOutput {
                config: cfg,
                rows: landing_rows.clone(),
            })
            .expect("serialize json");
            write_text_file(&path, &json);
            eprintln!("Wrote {}", path.display());
        }
        if let Some(path) = out_csv_landing {
            let csv = render_csv(&landing_rows);
            write_text_file(&path, &csv);
            eprintln!("Wrote {}", path.display());
        }
    }
}
