//! 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, lookup, and batched (`get_many`) lookup ns/op. Keys are generated by a
//! seeded RNG rather than materialized (a 500M-key vector would be ~16 GB). The lookup phases
//! re-seed 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).
//!
//! A `batch=<size>` arg sets the `get_many` batch size (default [DEFAULT_BATCH]). Each batch's keys
//! are sorted before the call, as a caller iterating sorted keys passes them.

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

#[global_allocator]
static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc::new(mimalloc::Config::DEFAULT);

// 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];

// Default `get_many` batch size. 10,000 matches the updates per commit in the QMDB `scale`
// benchmark's `generate`, so it approximates the keys one merkleize probes.
const DEFAULT_BATCH: usize = 10_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::detached(metric)
    }

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

/// Time three phases over the same `items` keys, returning the insert, `get` and `get_many`
/// durations:
///
/// - insert, which builds `index` one key at a time
/// - `get`, which looks every key up one at a time
/// - `get_many`, which looks every key up in sorted batches of `batch` keys and times only the calls
///
/// Both lookup phases read every value they find, as a caller would. Keys come from a seeded RNG
/// and are not stored, and each 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,
    batch: usize,
) -> (Duration, 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 mut sum = 0u64;
    let start = Instant::now();
    for _ in 0..items {
        for value in index.get(&rng.next_u64().to_be_bytes()) {
            sum = sum.wrapping_add(*value);
        }
    }
    let lookup = start.elapsed();
    black_box(sum);

    let mut rng = TestRng::new(SEED);
    let mut keys = Vec::with_capacity(batch);
    let mut many = Duration::ZERO;
    let mut remaining = items;
    while remaining > 0 {
        let n = remaining.min(Widen::widen(batch));
        keys.clear();
        keys.extend((0..n).map(|_| rng.next_u64().to_be_bytes()));
        keys.sort_unstable();
        let mut sum = 0u64;
        let start = Instant::now();
        index.get_many(&keys, |_, value| sum = sum.wrapping_add(*value));
        many += start.elapsed();
        black_box(sum);
        remaining -= n;
    }

    (insert, lookup, many)
}

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 batch =
        argv.iter()
            .find_map(|a| a.strip_prefix("batch="))
            .map_or(DEFAULT_BATCH, |b| match b.parse() {
                Ok(batch) if batch > 0 => batch,
                _ => panic!("batch size must be a positive integer, got {b:?}"),
            });
    let only: Vec<String> = argv
        .into_iter()
        .filter(|a| a.parse::<u64>().is_err() && !a.starts_with('-') && !a.starts_with("batch="))
        .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, with byte-array
        // keys and then with the integer keys QMDB's P=3 ordered config uses. 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, many) = measure($index, items, batch);
                    println!(
                        "  {:<32} insert={} ns/op  lookup={} ns/op  get_many(batch={batch})={} ns/op",
                        $name,
                        insert.as_nanos() / u128::from(items),
                        lookup.as_nanos() / u128::from(items),
                        many.as_nanos() / u128::from(items),
                    );
                }
            }};
        }

        run!(
            "partitioned_ordered_3",
            partitioned::ordered::Index::<_, _, 3>::new(DummyMetrics, Cap::<5>::new())
        );
        run!(
            "partitioned_ordered_3_eightcap",
            partitioned::ordered::Index::<_, _, 3>::new(DummyMetrics, EightCap)
        );
        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));
    }
}