//! Standalone, opt-in characterization of the index at huge key counts.
//!
//! At P=3 (16.8M partitions) criterion is a poor fit: it reruns each benchmark for many samples,
//! and every insert sample would rebuild a multi-GB index. This binary instead builds each variant
//! once and reports insert and lookup ns/op. Keys are generated by a seeded RNG rather than
//! materialized (a 500M-key vector would be ~16 GB); the lookup phase re-seeds to replay the same
//! keys, so the index itself is the only large allocation.
//!
//! Sizes are taken from numeric CLI args, e.g.
//! `cargo bench -p commonware-storage --bench index_scale -- 500000000`. With no args it runs the
//! default 20M/100M tier, but only when built with `--cfg huge_bench` -- so a bare `cargo bench`
//! (including CI's full-suite run) does nothing. Non-numeric args restrict which variants run by
//! name, e.g. `-- 1000000000 partitioned_ordered_3` runs only the SoA index at 1B (the heavy flat
//! baselines would need ~40+ GB at that size).

use commonware_runtime::{
    telemetry::metrics::{Metric, Registered, Registration},
    Metrics, Name, Supervisor,
};
use commonware_storage::{
    index::{ordered, partitioned, unordered, Unordered},
    translator::{Cap, EightCap},
};
use commonware_utils::TestRng;
use rand::Rng;
use std::{
    hint::black_box,
    time::{Duration, Instant},
};

// Default no-arg tier: 20M keys gives ~1.2 entries per P=3 partition, 100M gives ~6 -- enough to
// exercise the per-partition sorted runs.
const DEFAULT_ITEMS: [u64; 2] = [20_000_000, 100_000_000];

// Fixed RNG seed so the insert and lookup phases generate the same key sequence.
const SEED: u64 = 0;

/// No-op metrics context. Mirrors the criterion benches' helper; duplicated because a separate
/// `harness = false` target cannot share the criterion entry point's module.
#[derive(Clone)]
struct DummyMetrics;

impl Supervisor for DummyMetrics {
    fn child(&self, _: &'static str) -> Self {
        Self
    }

    fn with_attribute(self, _: &'static str, _: impl std::fmt::Display) -> Self {
        Self
    }

    fn name(&self) -> Name {
        Name::default()
    }
}

impl Metrics for DummyMetrics {
    fn register<N: Into<String>, H: Into<String>, M: Metric>(
        &self,
        _: N,
        _: H,
        metric: M,
    ) -> Registered<M> {
        Registered::with_registration(metric, Registration::from(()))
    }

    fn encode(&self) -> String {
        String::new()
    }
}

/// Insert `items` keys (timing the build), then look every key up (timing the gets), returning the
/// two batch durations. Keys come from a seeded RNG and are not stored; the lookup phase re-seeds to
/// replay the identical sequence. A fast RNG keeps per-key generation negligible next to the index
/// op, so the numbers stay comparable to materialized keys.
fn measure<I: Unordered<Value = u64>>(mut index: I, items: u64) -> (Duration, Duration) {
    let mut rng = TestRng::new(SEED);
    let start = Instant::now();
    for value in 0..items {
        index.insert(&rng.next_u64().to_be_bytes(), value);
    }
    let insert = start.elapsed();

    let mut rng = TestRng::new(SEED);
    let start = Instant::now();
    for _ in 0..items {
        black_box(index.get(&rng.next_u64().to_be_bytes()).next().is_some());
    }
    let lookup = start.elapsed();

    (insert, lookup)
}

fn main() {
    // Sizes come from numeric CLI args (e.g. `-- 500000000`). With none, use the default tier when
    // built with `--cfg huge_bench`; otherwise no-op, so a bare `cargo bench` (including CI's
    // full-suite run, which uses `--cfg full_bench`) does nothing.
    // Numeric args are sizes; other args name the variants to run (default: all). Flag-style args
    // (e.g. the `--bench` that `cargo bench` injects into the harness) are ignored.
    let argv: Vec<String> = std::env::args().skip(1).collect();
    let args: Vec<u64> = argv.iter().filter_map(|a| a.parse().ok()).collect();
    let only: Vec<String> = argv
        .into_iter()
        .filter(|a| a.parse::<u64>().is_err() && !a.starts_with('-'))
        .collect();
    let sizes = if !args.is_empty() {
        args
    } else if cfg!(huge_bench) {
        DEFAULT_ITEMS.to_vec()
    } else {
        eprintln!(
            "index_scale is opt-in; pass key counts (e.g. `-- 500000000`), or build with \
             `--cfg huge_bench` for the default 20M/100M tier"
        );
        return;
    };

    for items in sizes {
        println!("index_scale: items={items}");

        // Each variant is built once; insert is the timed build, lookup re-seeds and reuses the
        // populated index. P=3 ordered SoA (the structure under test) runs first; the flat BTree is
        // last because it is by far the slowest/heaviest baseline. (Ordered P=1/P=2 and unordered
        // P=3 are omitted: pathological build, or 16.8M hashmaps.)
        macro_rules! run {
            ($name:literal, $index:expr) => {{
                if only.is_empty() || only.iter().any(|v| v.as_str() == $name) {
                    let (insert, lookup) = measure($index, items);
                    println!(
                        "  {:<24} insert={} ns/op  lookup={} ns/op",
                        $name,
                        insert.as_nanos() / items as u128,
                        lookup.as_nanos() / items as u128,
                    );
                }
            }};
        }

        run!(
            "partitioned_ordered_3",
            partitioned::ordered::Index::<_, _, 3>::new(DummyMetrics, Cap::<5>::new())
        );
        run!("unordered", unordered::Index::new(DummyMetrics, EightCap));
        run!(
            "partitioned_unordered_1",
            partitioned::unordered::Index::<_, _, 1>::new(DummyMetrics, Cap::<7>::new())
        );
        run!(
            "partitioned_unordered_2",
            partitioned::unordered::Index::<_, _, 2>::new(DummyMetrics, Cap::<6>::new())
        );
        run!("ordered", ordered::Index::new(DummyMetrics, EightCap));
    }
}