//! Shared infrastructure for QMDB benchmarks: constants, config builders, type aliases, dispatch
//! macros, and the common `gen_random_kv` helper.

use commonware_cryptography::{DigestOf, Hasher as _, Sha256};
use commonware_parallel::Rayon;
use commonware_runtime::{buffer::paged::CacheRef, tokio::Context, BufferPooler, Strategizer};
use commonware_storage::{
    journal::contiguous::{fixed::Config as FConfig, variable::Config as VConfig},
    merkle::{self, full::Config as MerkleConfig, Family},
    qmdb::{
        any::{
            ordered::{
                fixed::{partitioned::p256::Db as OFixP256, Db as OFixed},
                variable::Db as OVariable,
            },
            traits::{DbAny, UnmerkleizedBatch},
            unordered::{fixed::Db as UFixed, variable::Db as UVariable},
            FixedConfig as AnyFixedConfig, VariableConfig as AnyVariableConfig,
        },
        current::{
            ordered::{fixed::Db as OCFixed, variable::Db as OCVariable},
            unordered::{fixed::Db as UCFixed, variable::Db as UCVariable},
            FixedConfig as CurrentFixedConfig, VariableConfig as CurrentVariableConfig,
        },
        immutable::fixed::{Config as ImmutableFixedConfig, Db as IFixed},
        keyless::variable::{Config as KeylessConfig, Db as Keyless},
    },
    translator::EightCap,
};
use commonware_utils::{NZUsize, TestRng, NZU16, NZU64};
use rand::{distr::Distribution, Rng};
use rand_distr::Zipf;
use std::num::{NonZeroU16, NonZeroU64, NonZeroUsize};

pub type Digest = DigestOf<Sha256>;

/// Default items per blob for benchmarks. This is small enough that blob boundary crossings can
/// affect benchmark time. Benchmarks that don't want to measure that cost should override via the
/// `_with` config generators.
pub const ITEMS_PER_BLOB: NonZeroU64 = NZU64!(50_000);
pub const CHUNK_SIZE: usize = 32;
pub const THREADS: NonZeroUsize = NZUsize!(8);
pub const PAGE_SIZE: NonZeroU16 = NZU16!(16384);
pub const PAGE_CACHE_SIZE: NonZeroUsize = NZUsize!(512);
pub const DELETE_FREQUENCY: u32 = 10;
pub const VARIABLE_VALUE_MAX_LEN: usize = 256;
pub const WRITE_BUFFER_SIZE: NonZeroUsize = NZUsize!(2 * 1024 * 1024);
pub const INIT_CACHE_SIZE: Option<NonZeroUsize> = Some(NZUsize!(1 << 18));

// -- Fixed value (Digest), fixed storage layout --

pub type AnyUFixDb<F> = UFixed<F, Context, Digest, Digest, Sha256, EightCap, Rayon>;
pub type AnyOFixDb<F> = OFixed<F, Context, Digest, Digest, Sha256, EightCap, Rayon>;
/// Ordered "any" DB with a partitioned snapshot index (256 partitions, P=1). Exercises the
/// partitioned ordered index's cursor (get_mut/find/update) on apply.
pub type AnyOFixP256Db<F> = OFixP256<F, Context, Digest, Digest, Sha256, EightCap, Rayon>;
pub type CurUFixDb<F> = UCFixed<F, Context, Digest, Digest, Sha256, EightCap, CHUNK_SIZE, Rayon>;
pub type CurOFixDb<F> = OCFixed<F, Context, Digest, Digest, Sha256, EightCap, CHUNK_SIZE, Rayon>;

// -- Fixed value (Digest), variable storage layout --
// Measures overhead of variable-capable storage when values are fixed-size.

pub type AnyUVarDigestDb<F> = UVariable<F, Context, Digest, Digest, Sha256, EightCap, Rayon>;
pub type AnyOVarDigestDb<F> = OVariable<F, Context, Digest, Digest, Sha256, EightCap, Rayon>;
pub type CurUVarDigestDb<F> =
    UCVariable<F, Context, Digest, Digest, Sha256, EightCap, CHUNK_SIZE, Rayon>;
pub type CurOVarDigestDb<F> =
    OCVariable<F, Context, Digest, Digest, Sha256, EightCap, CHUNK_SIZE, Rayon>;

// -- Variable value (Vec<u8>), variable storage layout --

pub type AnyUVarVecDb<F> = UVariable<F, Context, Digest, Vec<u8>, Sha256, EightCap, Rayon>;
pub type AnyOVarVecDb<F> = OVariable<F, Context, Digest, Vec<u8>, Sha256, EightCap, Rayon>;
pub type CurUVarVecDb<F> =
    UCVariable<F, Context, Digest, Vec<u8>, Sha256, EightCap, CHUNK_SIZE, Rayon>;
pub type CurOVarVecDb<F> =
    OCVariable<F, Context, Digest, Vec<u8>, Sha256, EightCap, CHUNK_SIZE, Rayon>;

// -- Immutable --

pub type ImmFixDb<F> = IFixed<F, Context, Digest, Digest, Sha256, EightCap, Rayon>;

// -- Keyless --

pub type KeylessDb<F> = Keyless<F, Context, Vec<u8>, Sha256, Rayon>;

/// Open a keyless benchmark database using the shared benchmark configuration.
pub async fn open_keyless_db<F: Family>(ctx: Context) -> KeylessDb<F> {
    let cfg = keyless_cfg(&ctx);
    KeylessDb::<F>::init(ctx, cfg).await.unwrap()
}

// -- Config builders --

const PARTITION_FIX: &str = "bench-fixed";
const PARTITION_VAR: &str = "bench-variable";
const PARTITION_KEYLESS: &str = "bench-keyless";
const PARTITION_IMM: &str = "bench-immutable";

fn merkle_cfg(
    suffix: &str,
    ctx: &(impl BufferPooler + Strategizer),
    page_cache: CacheRef,
    items_per_blob: NonZeroU64,
) -> MerkleConfig<Rayon> {
    MerkleConfig {
        journal_partition: format!("journal-{suffix}"),
        metadata_partition: format!("metadata-{suffix}"),
        items_per_blob,
        write_buffer: WRITE_BUFFER_SIZE,
        strategy: ctx.strategy(THREADS),
        page_cache,
    }
}

fn fix_log_cfg(suffix: &str, page_cache: CacheRef, items_per_blob: NonZeroU64) -> FConfig {
    FConfig {
        partition: format!("log-journal-{suffix}"),
        items_per_blob,
        page_cache,
        write_buffer: WRITE_BUFFER_SIZE,
    }
}

fn var_log_cfg<C>(
    suffix: &str,
    page_cache: CacheRef,
    codec_config: C,
    items_per_section: NonZeroU64,
) -> VConfig<C> {
    VConfig {
        partition: format!("log-journal-{suffix}"),
        items_per_section,
        compression: None,
        codec_config,
        page_cache,
        write_buffer: WRITE_BUFFER_SIZE,
    }
}

pub fn any_fix_cfg(ctx: &(impl BufferPooler + Strategizer)) -> AnyFixedConfig<EightCap, Rayon> {
    any_fix_cfg_with(ctx, ITEMS_PER_BLOB, PAGE_CACHE_SIZE)
}

pub fn any_fix_cfg_with(
    ctx: &(impl BufferPooler + Strategizer),
    items_per_blob: NonZeroU64,
    page_cache_size: NonZeroUsize,
) -> AnyFixedConfig<EightCap, Rayon> {
    let page_cache = CacheRef::from_pooler(ctx, PAGE_SIZE, page_cache_size);
    AnyFixedConfig {
        merkle_config: merkle_cfg(PARTITION_FIX, ctx, page_cache.clone(), items_per_blob),
        journal_config: fix_log_cfg(PARTITION_FIX, page_cache, items_per_blob),
        translator: EightCap,
        init_cache_size: INIT_CACHE_SIZE,
    }
}

pub fn imm_fix_cfg_with(
    ctx: &(impl BufferPooler + Strategizer),
    items_per_blob: NonZeroU64,
) -> ImmutableFixedConfig<EightCap, Rayon> {
    let page_cache = CacheRef::from_pooler(ctx, PAGE_SIZE, PAGE_CACHE_SIZE);
    ImmutableFixedConfig {
        merkle_config: merkle_cfg(PARTITION_IMM, ctx, page_cache.clone(), items_per_blob),
        log: fix_log_cfg(PARTITION_IMM, page_cache, items_per_blob),
        translator: EightCap,
        init_cache_size: INIT_CACHE_SIZE,
    }
}

pub fn cur_fix_cfg(ctx: &(impl BufferPooler + Strategizer)) -> CurrentFixedConfig<EightCap, Rayon> {
    cur_fix_cfg_with(ctx, ITEMS_PER_BLOB)
}

pub fn cur_fix_cfg_with(
    ctx: &(impl BufferPooler + Strategizer),
    items_per_blob: NonZeroU64,
) -> CurrentFixedConfig<EightCap, Rayon> {
    let page_cache = CacheRef::from_pooler(ctx, PAGE_SIZE, PAGE_CACHE_SIZE);
    CurrentFixedConfig {
        merkle_config: merkle_cfg(PARTITION_FIX, ctx, page_cache.clone(), items_per_blob),
        journal_config: fix_log_cfg(PARTITION_FIX, page_cache, items_per_blob),
        grafted_metadata_partition: format!("grafted-metadata-{PARTITION_FIX}"),
        translator: EightCap,
        init_cache_size: INIT_CACHE_SIZE,
    }
}

pub fn any_var_digest_cfg(
    ctx: &(impl BufferPooler + Strategizer),
) -> AnyVariableConfig<EightCap, ((), ()), Rayon> {
    any_var_digest_cfg_with(ctx, ITEMS_PER_BLOB)
}

pub fn any_var_digest_cfg_with(
    ctx: &(impl BufferPooler + Strategizer),
    items_per_blob: NonZeroU64,
) -> AnyVariableConfig<EightCap, ((), ()), Rayon> {
    let page_cache = CacheRef::from_pooler(ctx, PAGE_SIZE, PAGE_CACHE_SIZE);
    AnyVariableConfig {
        merkle_config: merkle_cfg(PARTITION_VAR, ctx, page_cache.clone(), items_per_blob),
        journal_config: var_log_cfg(PARTITION_VAR, page_cache, ((), ()), items_per_blob),
        translator: EightCap,
        init_cache_size: INIT_CACHE_SIZE,
    }
}

pub fn cur_var_digest_cfg(
    ctx: &(impl BufferPooler + Strategizer),
) -> CurrentVariableConfig<EightCap, ((), ()), Rayon> {
    cur_var_digest_cfg_with(ctx, ITEMS_PER_BLOB)
}

pub fn cur_var_digest_cfg_with(
    ctx: &(impl BufferPooler + Strategizer),
    items_per_blob: NonZeroU64,
) -> CurrentVariableConfig<EightCap, ((), ()), Rayon> {
    let page_cache = CacheRef::from_pooler(ctx, PAGE_SIZE, PAGE_CACHE_SIZE);
    CurrentVariableConfig {
        merkle_config: merkle_cfg(PARTITION_VAR, ctx, page_cache.clone(), items_per_blob),
        journal_config: var_log_cfg(PARTITION_VAR, page_cache, ((), ()), items_per_blob),
        grafted_metadata_partition: format!("grafted-metadata-{PARTITION_VAR}"),
        translator: EightCap,
        init_cache_size: INIT_CACHE_SIZE,
    }
}

/// Codec config for variable-length `Vec<u8>` values.
type VarVecCfg = ((), (commonware_codec::RangeCfg<usize>, ()));

pub fn any_var_vec_cfg(
    ctx: &(impl BufferPooler + Strategizer),
) -> AnyVariableConfig<EightCap, VarVecCfg, Rayon> {
    any_var_vec_cfg_with(ctx, ITEMS_PER_BLOB)
}

pub fn any_var_vec_cfg_with(
    ctx: &(impl BufferPooler + Strategizer),
    items_per_blob: NonZeroU64,
) -> AnyVariableConfig<EightCap, VarVecCfg, Rayon> {
    let page_cache = CacheRef::from_pooler(ctx, PAGE_SIZE, PAGE_CACHE_SIZE);
    AnyVariableConfig {
        merkle_config: merkle_cfg(PARTITION_VAR, ctx, page_cache.clone(), items_per_blob),
        journal_config: var_log_cfg(
            PARTITION_VAR,
            page_cache,
            ((), ((0..=10000).into(), ())),
            items_per_blob,
        ),
        translator: EightCap,
        init_cache_size: INIT_CACHE_SIZE,
    }
}

pub fn cur_var_vec_cfg(
    ctx: &(impl BufferPooler + Strategizer),
) -> CurrentVariableConfig<EightCap, VarVecCfg, Rayon> {
    cur_var_vec_cfg_with(ctx, ITEMS_PER_BLOB)
}

pub fn cur_var_vec_cfg_with(
    ctx: &(impl BufferPooler + Strategizer),
    items_per_blob: NonZeroU64,
) -> CurrentVariableConfig<EightCap, VarVecCfg, Rayon> {
    let page_cache = CacheRef::from_pooler(ctx, PAGE_SIZE, PAGE_CACHE_SIZE);
    CurrentVariableConfig {
        merkle_config: merkle_cfg(PARTITION_VAR, ctx, page_cache.clone(), items_per_blob),
        journal_config: var_log_cfg(
            PARTITION_VAR,
            page_cache,
            ((), ((0..=10000).into(), ())),
            items_per_blob,
        ),
        grafted_metadata_partition: format!("grafted-metadata-{PARTITION_VAR}"),
        translator: EightCap,
        init_cache_size: INIT_CACHE_SIZE,
    }
}

pub fn keyless_cfg(
    ctx: &(impl BufferPooler + Strategizer),
) -> KeylessConfig<(commonware_codec::RangeCfg<usize>, ()), Rayon> {
    keyless_cfg_with(ctx, ITEMS_PER_BLOB)
}

pub fn keyless_cfg_with(
    ctx: &(impl BufferPooler + Strategizer),
    items_per_blob: NonZeroU64,
) -> KeylessConfig<(commonware_codec::RangeCfg<usize>, ()), Rayon> {
    let page_cache = CacheRef::from_pooler(ctx, PAGE_SIZE, PAGE_CACHE_SIZE);
    KeylessConfig {
        merkle: merkle_cfg(PARTITION_KEYLESS, ctx, page_cache.clone(), items_per_blob),
        log: var_log_cfg(
            PARTITION_KEYLESS,
            page_cache,
            ((0..=10000).into(), ()),
            items_per_blob,
        ),
    }
}

// -- Shared variant definitions --

macro_rules! define_db_variants {
    (
        enum $enum_name:ident;
        const $variants_name:ident;
        dispatch $dispatch_name:ident;
        timed_dispatch $timed_dispatch_name:ident;
        entries = [
            $(
                {
                    entry: $entry:ident,
                    name: $name:literal,
                    db: $db:ty,
                    cfg: $cfg:path,
                }
            )+
        ];
    ) => {
        #[derive(Debug, Clone, Copy)]
        enum $enum_name {
            $($entry),+
        }

        impl $enum_name {
            const fn name(self) -> &'static str {
                match self {
                    $(Self::$entry => $name),+
                }
            }
        }

        const $variants_name: &[$enum_name] = &[$($enum_name::$entry),+];

        macro_rules! $dispatch_name {
            ($ctx_expr:expr, $variant_expr:expr, |$db_name:ident| $body:expr) => {
                match $variant_expr {
                    $(
                        $enum_name::$entry => {
                            let ctx = $ctx_expr;
                            let cfg = $cfg(&ctx);
                            #[allow(unused_mut)]
                            let mut $db_name = <$db>::init(ctx.child("storage"), cfg).await.unwrap();
                            $body
                        }
                    )+
                }
            };
        }

        #[allow(unused_macros)]
        macro_rules! $timed_dispatch_name {
            ($ctx_expr:expr, $variant_expr:expr, $iters:expr, $cache_size:expr, |$db_name:ident| $body:expr) => {
                match $variant_expr {
                    $(
                        $enum_name::$entry => {
                            let ctx = $ctx_expr;
                            let mut cfg = $cfg(&ctx);
                            cfg.init_cache_size = $cache_size;
                            let start = std::time::Instant::now();
                            for _ in 0..$iters {
                                #[allow(unused_mut)]
                                let mut $db_name =
                                    <$db>::init(ctx.child("storage"), cfg.clone()).await.unwrap();
                                $body
                            }
                            start.elapsed()
                        }
                    )+
                }
            };
        }
    };
}

pub(crate) use define_db_variants;

macro_rules! define_fixed_variants {
    (
        enum $enum_name:ident;
        const $variants_name:ident;
        dispatch $dispatch_name:ident;
        timed_dispatch $timed_dispatch_name:ident;
    ) => {
        $crate::common::define_db_variants! {
            enum $enum_name;
            const $variants_name;
            dispatch $dispatch_name;
            timed_dispatch $timed_dispatch_name;
            entries = [
                {
                    entry: AnyUnorderedFixedMmr,
                    name: "any::unordered::fixed::mmr",
                    db: $crate::common::AnyUFixDb<commonware_storage::merkle::mmr::Family>,
                    cfg: $crate::common::any_fix_cfg,
                }
                {
                    entry: AnyUnorderedFixedMmb,
                    name: "any::unordered::fixed::mmb",
                    db: $crate::common::AnyUFixDb<commonware_storage::merkle::mmb::Family>,
                    cfg: $crate::common::any_fix_cfg,
                }
                {
                    entry: AnyOrderedFixedMmr,
                    name: "any::ordered::fixed::mmr",
                    db: $crate::common::AnyOFixDb<commonware_storage::merkle::mmr::Family>,
                    cfg: $crate::common::any_fix_cfg,
                }
                {
                    entry: AnyOrderedFixedMmb,
                    name: "any::ordered::fixed::mmb",
                    db: $crate::common::AnyOFixDb<commonware_storage::merkle::mmb::Family>,
                    cfg: $crate::common::any_fix_cfg,
                }
                {
                    entry: AnyUnorderedVariableMmr,
                    name: "any::unordered::variable::mmr",
                    db: $crate::common::AnyUVarDigestDb<commonware_storage::merkle::mmr::Family>,
                    cfg: $crate::common::any_var_digest_cfg,
                }
                {
                    entry: AnyUnorderedVariableMmb,
                    name: "any::unordered::variable::mmb",
                    db: $crate::common::AnyUVarDigestDb<commonware_storage::merkle::mmb::Family>,
                    cfg: $crate::common::any_var_digest_cfg,
                }
                {
                    entry: AnyOrderedVariableMmr,
                    name: "any::ordered::variable::mmr",
                    db: $crate::common::AnyOVarDigestDb<commonware_storage::merkle::mmr::Family>,
                    cfg: $crate::common::any_var_digest_cfg,
                }
                {
                    entry: AnyOrderedVariableMmb,
                    name: "any::ordered::variable::mmb",
                    db: $crate::common::AnyOVarDigestDb<commonware_storage::merkle::mmb::Family>,
                    cfg: $crate::common::any_var_digest_cfg,
                }
                {
                    entry: CurrentUnorderedFixedMmr,
                    name: "current::unordered::fixed::mmr",
                    db: $crate::common::CurUFixDb<commonware_storage::merkle::mmr::Family>,
                    cfg: $crate::common::cur_fix_cfg,
                }
                {
                    entry: CurrentUnorderedFixedMmb,
                    name: "current::unordered::fixed::mmb",
                    db: $crate::common::CurUFixDb<commonware_storage::merkle::mmb::Family>,
                    cfg: $crate::common::cur_fix_cfg,
                }
                {
                    entry: CurrentOrderedFixedMmr,
                    name: "current::ordered::fixed::mmr",
                    db: $crate::common::CurOFixDb<commonware_storage::merkle::mmr::Family>,
                    cfg: $crate::common::cur_fix_cfg,
                }
                {
                    entry: CurrentOrderedFixedMmb,
                    name: "current::ordered::fixed::mmb",
                    db: $crate::common::CurOFixDb<commonware_storage::merkle::mmb::Family>,
                    cfg: $crate::common::cur_fix_cfg,
                }
                {
                    entry: CurrentUnorderedVariableMmr,
                    name: "current::unordered::variable::mmr",
                    db: $crate::common::CurUVarDigestDb<commonware_storage::merkle::mmr::Family>,
                    cfg: $crate::common::cur_var_digest_cfg,
                }
                {
                    entry: CurrentUnorderedVariableMmb,
                    name: "current::unordered::variable::mmb",
                    db: $crate::common::CurUVarDigestDb<commonware_storage::merkle::mmb::Family>,
                    cfg: $crate::common::cur_var_digest_cfg,
                }
                {
                    entry: CurrentOrderedVariableMmr,
                    name: "current::ordered::variable::mmr",
                    db: $crate::common::CurOVarDigestDb<commonware_storage::merkle::mmr::Family>,
                    cfg: $crate::common::cur_var_digest_cfg,
                }
                {
                    entry: CurrentOrderedVariableMmb,
                    name: "current::ordered::variable::mmb",
                    db: $crate::common::CurOVarDigestDb<commonware_storage::merkle::mmb::Family>,
                    cfg: $crate::common::cur_var_digest_cfg,
                }
            ];
        }
    };
}

pub(crate) use define_fixed_variants;

macro_rules! define_vec_variants {
    (
        enum $enum_name:ident;
        const $variants_name:ident;
        dispatch $dispatch_name:ident;
        timed_dispatch $timed_dispatch_name:ident;
    ) => {
        $crate::common::define_db_variants! {
            enum $enum_name;
            const $variants_name;
            dispatch $dispatch_name;
            timed_dispatch $timed_dispatch_name;
            entries = [
                {
                    entry: AnyUnorderedMmr,
                    name: "any::unordered::variable-vec::mmr",
                    db: $crate::common::AnyUVarVecDb<commonware_storage::merkle::mmr::Family>,
                    cfg: $crate::common::any_var_vec_cfg,
                }
                {
                    entry: AnyUnorderedMmb,
                    name: "any::unordered::variable-vec::mmb",
                    db: $crate::common::AnyUVarVecDb<commonware_storage::merkle::mmb::Family>,
                    cfg: $crate::common::any_var_vec_cfg,
                }
                {
                    entry: AnyOrderedMmr,
                    name: "any::ordered::variable-vec::mmr",
                    db: $crate::common::AnyOVarVecDb<commonware_storage::merkle::mmr::Family>,
                    cfg: $crate::common::any_var_vec_cfg,
                }
                {
                    entry: AnyOrderedMmb,
                    name: "any::ordered::variable-vec::mmb",
                    db: $crate::common::AnyOVarVecDb<commonware_storage::merkle::mmb::Family>,
                    cfg: $crate::common::any_var_vec_cfg,
                }
                {
                    entry: CurrentUnorderedMmr,
                    name: "current::unordered::variable-vec::mmr",
                    db: $crate::common::CurUVarVecDb<commonware_storage::merkle::mmr::Family>,
                    cfg: $crate::common::cur_var_vec_cfg,
                }
                {
                    entry: CurrentUnorderedMmb,
                    name: "current::unordered::variable-vec::mmb",
                    db: $crate::common::CurUVarVecDb<commonware_storage::merkle::mmb::Family>,
                    cfg: $crate::common::cur_var_vec_cfg,
                }
                {
                    entry: CurrentOrderedMmr,
                    name: "current::ordered::variable-vec::mmr",
                    db: $crate::common::CurOVarVecDb<commonware_storage::merkle::mmr::Family>,
                    cfg: $crate::common::cur_var_vec_cfg,
                }
                {
                    entry: CurrentOrderedMmb,
                    name: "current::ordered::variable-vec::mmb",
                    db: $crate::common::CurOVarVecDb<commonware_storage::merkle::mmb::Family>,
                    cfg: $crate::common::cur_var_vec_cfg,
                }
            ];
        }
    };
}

pub(crate) use define_vec_variants;

// -- Data generation --

/// Seed a database with `num_elements` entries, then perform `num_operations` random
/// updates/deletes. Commits periodically when `commit_frequency` is `Some`.
///
/// `seed_batch` caps how many seeds accumulate before each merkleize+apply+commit (bounding peak
/// memory); `None` seeds in a single batch. `prune_frequency` prunes to the inactivity floor every
/// `Some(n)` commits to bound on-disk growth; `None` never prunes during generation.
///
/// `key_zipf_exponent` selects how update/delete keys are sampled: `None` is uniform; `Some(s)`
/// draws keys from a Zipf distribution with exponent `s`, so a hot subset is churned far more than
/// the long tail (a more realistic workload than uniform churn).
///
/// `keyspace` sets the index space that updates are drawn from; the seed is always the distinct keys
/// `0..num_elements`. `None` means updates also draw from `0..num_elements`, so they only touch seeded
/// keys. `Some(k)` (with `k >= num_elements`) draws updates over all of `0..k`, so some updates land on
/// unseeded keys and insert them organically -- a growing keyspace rather than a fixed population.
#[allow(clippy::too_many_arguments)]
pub async fn gen_random_kv<F, M>(
    db: &mut M,
    num_elements: u64,
    num_operations: u64,
    commit_frequency: Option<u32>,
    seed_batch: Option<u64>,
    prune_frequency: Option<u32>,
    key_zipf_exponent: Option<f64>,
    keyspace: Option<u64>,
    make_value: impl Fn(&mut TestRng) -> M::Value,
) where
    F: Family,
    M: DbAny<F, Key = Digest>,
{
    let mut rng = TestRng::new(42);

    // Count commits across both phases so `prune_frequency` can prune to the inactivity floor every
    // N commits, bounding on-disk growth instead of accumulating the whole (re-appended) log until
    // the end. Pruning below the floor never affects an `init` replay, which starts at the floor.
    let mut commits = 0u32;

    // Seed the db with `num_elements` entries. `seed_batch` caps how many seeds accumulate before a
    // merkleize+apply (bounding the batch); `None` seeds in a single batch. Each apply is followed by
    // a `commit`, whose `merkle.flush()` writes the in-memory merkle nodes to the journal and prunes
    // them from memory -- without that, a large build accumulates the whole merkle in RAM.
    {
        let mut batch = db.new_batch();
        let mut pending = 0u64;
        for i in 0u64..num_elements {
            let key = Sha256::hash(&i.to_be_bytes());
            batch = batch.write(key, Some(make_value(&mut rng)));
            pending += 1;
            if seed_batch.is_some_and(|n| pending >= n) {
                let merkleized = batch.merkleize(db, None).await.unwrap();
                db.apply_batch(merkleized).await.unwrap();
                db.commit().await.unwrap();
                commits += 1;
                if prune_frequency.is_some_and(|n| commits.is_multiple_of(n)) {
                    let boundary = db.sync_boundary();
                    db.prune(boundary).await.unwrap();
                }
                batch = db.new_batch();
                pending = 0;
            }
        }
        if pending > 0 {
            let merkleized = batch.merkleize(db, None).await.unwrap();
            db.apply_batch(merkleized).await.unwrap();
            db.commit().await.unwrap();
        }
    }

    // Perform `num_operations` random updates/deletes, committing periodically. Each apply is
    // followed by a `commit` so the merkle structure is flushed out of memory (see the seed phase).
    {
        // Sample over the full keyspace (which may exceed the seeded set, so some samples hit unseeded
        // keys and insert them). `Zipf::new` samples a rank in `[1, space]`; map it to a 0-based index.
        let space = keyspace.unwrap_or(num_elements);
        let zipf =
            key_zipf_exponent.map(|s| Zipf::new(space as f64, s).expect("valid zipf parameters"));
        let mut batch = db.new_batch();
        for _ in 0u64..num_operations {
            let idx = match &zipf {
                Some(z) => ((z.sample(&mut rng) as u64).saturating_sub(1)).min(space - 1),
                None => rng.next_u64() % space,
            };
            let rand_key = Sha256::hash(&idx.to_be_bytes());
            if rng.next_u32().is_multiple_of(DELETE_FREQUENCY) {
                batch = batch.write(rand_key, None);
                continue;
            }
            batch = batch.write(rand_key, Some(make_value(&mut rng)));
            if let Some(freq) = commit_frequency {
                if rng.next_u32().is_multiple_of(freq) {
                    let merkleized = batch.merkleize(db, None).await.unwrap();
                    db.apply_batch(merkleized).await.unwrap();
                    db.commit().await.unwrap();
                    commits += 1;
                    if prune_frequency.is_some_and(|n| commits.is_multiple_of(n)) {
                        let boundary = db.sync_boundary();
                        db.prune(boundary).await.unwrap();
                    }
                    batch = db.new_batch();
                }
            }
        }
        let merkleized = batch.merkleize(db, None).await.unwrap();
        db.apply_batch(merkleized).await.unwrap();
        db.commit().await.unwrap();
    }
}

/// Generate a fixed-size digest value.
pub fn make_fixed_value(rng: &mut TestRng) -> Digest {
    Sha256::hash(&rng.next_u32().to_be_bytes())
}

/// Pre-populate the database with `num_keys` unique keys, then commit.
pub async fn seed_db<F: merkle::Family, C: DbAny<F, Key = Digest, Value = Digest>>(
    db: &mut C,
    num_keys: u64,
) {
    let mut rng = TestRng::new(42);
    let mut batch = db.new_batch();
    for i in 0u64..num_keys {
        let k = Sha256::hash(&i.to_be_bytes());
        batch = batch.write(k, Some(make_fixed_value(&mut rng)));
    }
    let merkleized = batch.merkleize(db, None).await.unwrap();
    db.apply_batch(merkleized).await.unwrap();
    db.commit().await.unwrap();
}

/// Write `num_updates` random key updates into a batch.
pub fn write_random_updates<B, Db>(
    mut batch: B,
    num_updates: u64,
    num_keys: u64,
    rng: &mut TestRng,
) -> B
where
    B: UnmerkleizedBatch<Db, K = Digest, V = Digest>,
    Db: ?Sized,
{
    for _ in 0..num_updates {
        let idx = rng.next_u64() % num_keys;
        let k = Sha256::hash(&idx.to_be_bytes());
        batch = batch.write(k, Some(make_fixed_value(rng)));
    }
    batch
}

/// Generate a variable-size `Vec<u8>` value (1-256 bytes).
pub fn make_var_value(rng: &mut TestRng) -> Vec<u8> {
    let len = (rng.next_u32() as usize) % VARIABLE_VALUE_MAX_LEN + 1;
    vec![rng.next_u32() as u8; len]
}