//! Benchmarks for QMDB startup initialization performance.
//!
//! These benchmarks have expensive setup (generating large random databases) that runs lazily
//! inside `bench_function` so criterion's name filter can skip them entirely.

use crate::common::{
    define_fixed_variants, define_vec_variants, gen_random_kv, make_fixed_value, make_var_value,
    Digest,
};
use commonware_runtime::{
    benchmarks::{context, tokio},
    tokio::{Config, Context},
    Runner as _, Supervisor as _,
};
use commonware_storage::{merkle::Family, qmdb::any::traits::DbAny};
use commonware_utils::{NZUsize, TestRng};
use core::num::NonZeroUsize;
use criterion::{criterion_group, Criterion};

const NUM_ELEMENTS: u64 = 100_000;
const NUM_OPERATIONS: u64 = 1_000_000;
const COMMIT_FREQUENCY: u32 = 10_000;

/// Init-time `(location -> key)` cache sizes to compare: `None` disables the cache (the no-cache
/// baseline), and a reasonably sized cache that covers the bench's working set.
const CACHE_SIZES: [Option<NonZeroUsize>; 2] = [None, Some(NZUsize!(1 << 18))];

cfg_if::cfg_if! {
    if #[cfg(not(full_bench))] {
        const CASES: [(u64, u64); 1] = [(NUM_ELEMENTS, NUM_OPERATIONS)];
    } else {
        const CASES: [(u64, u64); 2] = [
            (NUM_ELEMENTS, NUM_OPERATIONS),
            (NUM_ELEMENTS * 3, NUM_OPERATIONS * 3),
        ];
    }
}

/// Populate, prune, and sync a database (used in setup phase).
async fn populate_and_sync<F: Family, C: DbAny<F, Key = Digest>>(
    db: &mut C,
    elements: u64,
    operations: u64,
    make_value: impl Fn(&mut TestRng) -> C::Value,
) {
    gen_random_kv::<F, _>(
        db,
        elements,
        operations,
        Some(COMMIT_FREQUENCY),
        None, // seed_batch
        None, // prune_frequency
        None, // key_zipf_exponent (uniform churn)
        None, // keyspace (all keys seeded)
        make_value,
    )
    .await;
    db.prune(db.sync_boundary()).await.unwrap();
    db.sync().await.unwrap();
}

// -- Fixed-value variants (16 = 8 db shapes x 2 merkle families) --

define_fixed_variants! {
    enum FixedVariant;
    const FIXED_VARIANTS;
    dispatch dispatch_fixed;
    timed_dispatch dispatch_fixed_timed_init;
}

fn bench_fixed_value_init(c: &mut Criterion) {
    let cfg = Config::default();
    for (elements, operations) in CASES {
        for &variant in FIXED_VARIANTS {
            // Populated lazily on the first sample of the first matched cache size, then reused by
            // every cache size for this variant (all read the same on-disk database).
            let mut initialized = false;
            for &cache_size in &CACHE_SIZES {
                let cache = cache_size.map_or(0, NonZeroUsize::get);
                let runner = tokio::Runner::new(cfg.clone());
                c.bench_function(
                    &format!(
                        "{}/variant={} cache={cache} elements={elements}",
                        module_path!(),
                        variant.name(),
                    ),
                    |b| {
                        // Setup: populate database (once, on first matched sample).
                        if !initialized {
                            commonware_runtime::tokio::Runner::new(cfg.clone()).start(
                                |ctx| async move {
                                    dispatch_fixed!(ctx, variant, |db| {
                                        populate_and_sync(
                                            &mut db,
                                            elements,
                                            operations,
                                            make_fixed_value,
                                        )
                                        .await;
                                    });
                                },
                            );
                            initialized = true;
                        }

                        // Benchmark: measure init time at this cache size.
                        b.to_async(&runner).iter_custom(move |iters| async move {
                            let ctx = context::get::<Context>();
                            dispatch_fixed_timed_init!(ctx, variant, iters, cache_size, |db| {
                                assert_ne!(db.bounds().end, 0);
                            })
                        });
                    },
                );
            }

            // Cleanup: destroy database.
            if initialized {
                commonware_runtime::tokio::Runner::new(cfg.clone()).start(|ctx| async move {
                    dispatch_fixed!(ctx, variant, |db| {
                        db.destroy().await.unwrap();
                    });
                });
            }
        }
    }
}

// -- Variable-value variants (8 = 4 db shapes x 2 merkle families) --

define_vec_variants! {
    enum VarVariant;
    const VEC_VARIANTS;
    dispatch dispatch_var;
    timed_dispatch dispatch_var_timed_init;
}

fn bench_var_value_init(c: &mut Criterion) {
    let cfg = Config::default();
    for (elements, operations) in CASES {
        for &variant in VEC_VARIANTS {
            // Populated lazily on the first sample of the first matched cache size, then reused by
            // every cache size for this variant (all read the same on-disk database).
            let mut initialized = false;
            for &cache_size in &CACHE_SIZES {
                let cache = cache_size.map_or(0, NonZeroUsize::get);
                let runner = tokio::Runner::new(cfg.clone());
                c.bench_function(
                    &format!(
                        "{}/variant={} cache={cache} elements={elements}",
                        module_path!(),
                        variant.name(),
                    ),
                    |b| {
                        // Setup: populate database (once, on first matched sample).
                        if !initialized {
                            commonware_runtime::tokio::Runner::new(cfg.clone()).start(
                                |ctx| async move {
                                    dispatch_var!(ctx, variant, |db| {
                                        populate_and_sync(
                                            &mut db,
                                            elements,
                                            operations,
                                            make_var_value,
                                        )
                                        .await;
                                    });
                                },
                            );
                            initialized = true;
                        }

                        // Benchmark: measure init time at this cache size.
                        b.to_async(&runner).iter_custom(move |iters| async move {
                            let ctx = context::get::<Context>();
                            dispatch_var_timed_init!(ctx, variant, iters, cache_size, |db| {
                                assert_ne!(db.bounds().end, 0);
                            })
                        });
                    },
                );
            }

            // Cleanup: destroy database.
            if initialized {
                commonware_runtime::tokio::Runner::new(cfg.clone()).start(|ctx| async move {
                    dispatch_var!(ctx, variant, |db| {
                        db.destroy().await.unwrap();
                    });
                });
            }
        }
    }
}

criterion_group! {
    name = benches;
    config = Criterion::default().sample_size(10);
    targets = bench_fixed_value_init, bench_var_value_init
}