//! An Any database implementation with an unordered key space and fixed-size values.

use crate::{
    index::unordered::Index,
    journal::contiguous::fixed::Journal,
    merkle::{Family, Location},
    qmdb::{
        any::{unordered, value::FixedEncoding, FixedConfig as Config, FixedValue},
        Error,
    },
    translator::Translator,
    Context,
};
use commonware_cryptography::Hasher;
use commonware_parallel::Strategy;
use commonware_utils::Array;

pub type Update<K, V> = unordered::Update<K, FixedEncoding<V>>;
pub type Operation<F, K, V> = unordered::Operation<F, K, FixedEncoding<V>>;

/// A key-value QMDB based on an authenticated log of operations, supporting authentication of any
/// value ever associated with a key.
pub type Db<F, E, K, V, H, T, S> = super::Db<
    F,
    E,
    Journal<E, Operation<F, K, V>>,
    Index<T, Location<F>>,
    H,
    Update<K, V>,
    { crate::qmdb::any::BITMAP_CHUNK_BYTES },
    S,
>;

impl<F: Family, E: Context, K: Array, V: FixedValue, H: Hasher, T: Translator, S: Strategy>
    Db<F, E, K, V, H, T, S>
{
    /// Returns a [Db] QMDB initialized from `cfg`. Uncommitted log operations will be
    /// discarded and the state of the db will be as of the last committed operation.
    pub async fn init(context: E, cfg: Config<T, S>) -> Result<Self, Error<F>> {
        crate::qmdb::any::init(context, cfg).await
    }
}

/// Partitioned index variants that divide the key space into `2^(P*8)` partitions.
///
/// See [partitioned::Db] for the generic type, or use the convenience aliases:
/// - [partitioned::p256::Db] for 256 partitions (P=1)
/// - [partitioned::p64k::Db] for 65,536 partitions (P=2)
pub mod partitioned {
    pub use super::{Operation, Update};
    use crate::{
        index::partitioned::unordered::Index,
        journal::contiguous::fixed::Journal,
        merkle::{Family, Location},
        qmdb::{
            any::{FixedConfig as Config, FixedValue},
            Error,
        },
        translator::Translator,
        Context,
    };
    use commonware_cryptography::Hasher;
    use commonware_parallel::Strategy;
    use commonware_utils::Array;

    /// A key-value QMDB with a partitioned snapshot index.
    ///
    /// This is the partitioned variant of [super::Db]. The const generic `P` specifies
    /// the number of prefix bytes used for partitioning:
    /// - `P = 1`: 256 partitions
    /// - `P = 2`: 65,536 partitions
    ///
    /// Use partitioned indices when you have a large number of keys (>> 2^(P*8)) and memory
    /// efficiency is important. Keys should be uniformly distributed across the prefix space.
    pub type Db<F, E, K, V, H, T, const P: usize, S> = crate::qmdb::any::unordered::Db<
        F,
        E,
        Journal<E, Operation<F, K, V>>,
        Index<T, Location<F>, P>,
        H,
        Update<K, V>,
        { crate::qmdb::any::BITMAP_CHUNK_BYTES },
        S,
    >;

    impl<
            F: Family,
            E: Context,
            K: Array,
            V: FixedValue,
            H: Hasher,
            T: Translator,
            const P: usize,
            S: Strategy,
        > Db<F, E, K, V, H, T, P, S>
    {
        /// Returns a [Db] QMDB initialized from `cfg`. Uncommitted log operations will be
        /// discarded and the state of the db will be as of the last committed operation.
        pub async fn init(context: E, cfg: Config<T, S>) -> Result<Self, Error<F>> {
            crate::qmdb::any::init(context, cfg).await
        }
    }

    /// Convenience type aliases for 256 partitions (P=1).
    pub mod p256 {
        /// Fixed-value DB with 256 partitions.
        pub type Db<F, E, K, V, H, T, S> = super::Db<F, E, K, V, H, T, 1, S>;
    }

    /// Convenience type aliases for 65,536 partitions (P=2).
    pub mod p64k {
        /// Fixed-value DB with 65,536 partitions.
        pub type Db<F, E, K, V, H, T, S> = super::Db<F, E, K, V, H, T, 2, S>;
    }
}

// pub(crate) so helpers can be used by the sync module.
#[cfg(test)]
pub(crate) mod test {
    use super::*;
    use crate::{
        index::Unordered as _,
        merkle::{
            mmr::{self, Location},
            Location as GenericLocation,
        },
        qmdb::{
            any::{
                test::{fixed_db_config, fixed_db_config_with_strategy},
                unordered::{fixed::Operation, Update},
            },
            verify_proof,
        },
        translator::TwoCap,
    };
    use commonware_cryptography::{sha256::Digest, Sha256};
    use commonware_macros::{select, test_traced};
    use commonware_math::algebra::Random;
    use commonware_parallel::{Rayon, Sequential};
    use commonware_runtime::{
        deterministic::{self, Context},
        Clock as _, Metrics as _, Runner as _, Strategizer as _, Supervisor as _,
    };
    use commonware_utils::{NZUsize, TestRng, NZU64};
    use core::num::NonZeroUsize;
    use futures::FutureExt as _;
    use rand::Rng;
    use std::{collections::HashMap, time::Duration};

    /// A generic type alias for an Any database parameterized by merkle family.
    type AnyTestGeneric<F> = crate::qmdb::any::db::Db<
        F,
        deterministic::Context,
        Journal<
            deterministic::Context,
            crate::qmdb::any::operation::Unordered<F, Digest, FixedEncoding<Digest>>,
        >,
        Index<TwoCap, GenericLocation<F>>,
        Sha256,
        crate::qmdb::any::operation::update::Unordered<Digest, FixedEncoding<Digest>>,
        { crate::qmdb::any::BITMAP_CHUNK_BYTES },
        Sequential,
    >;

    /// A type alias for the concrete [Db] type used in these unit tests.
    pub(crate) type AnyTest =
        Db<mmr::Family, deterministic::Context, Digest, Digest, Sha256, TwoCap, Sequential>;

    /// Return an `Any` database initialized with a fixed config, generic over merkle family.
    async fn open_db_generic<F: Family>(context: deterministic::Context) -> AnyTestGeneric<F> {
        let cfg = fixed_db_config::<TwoCap>("partition", &context);
        crate::qmdb::any::init(context, cfg).await.unwrap()
    }

    /// Create a test database with unique partition names
    pub(crate) async fn create_test_db(mut context: Context) -> AnyTest {
        let seed = context.next_u64();
        let cfg = fixed_db_config::<TwoCap>(&seed.to_string(), &context);
        AnyTest::init(context, cfg).await.unwrap()
    }

    /// `get_many` over a batch large enough for the fused sharded path matches per-key `get`.
    #[test_traced]
    fn test_get_many_fused_sharded_matches_get() {
        let executor = deterministic::Runner::default();
        executor.start(|context| async move {
            // The fused path requires parallelism > 1 (the Sequential test config never takes
            // it) and at least 4096 keys. The tiny test page cache pushes most keys through
            // the batched miss fallback, and TwoCap produces translated-key collisions.
            type ParTest = Db<mmr::Family, Context, Digest, Digest, Sha256, TwoCap, Rayon>;
            let strategy = context.strategy(NZUsize!(2));
            let cfg = fixed_db_config_with_strategy::<TwoCap, Rayon>("fused", &context, strategy);
            let mut db = ParTest::init(context, cfg).await.unwrap();

            let mut rng = TestRng::new(7);
            let mut keys = Vec::with_capacity(4300);
            let mut batch = db.new_batch();
            for _ in 0..4200 {
                let key = Digest::random(&mut rng);
                let value = Digest::random(&mut rng);
                keys.push(key);
                batch = batch.write(key, Some(value));
            }
            let merkleized = batch.merkleize(&db, None).await.unwrap();
            db.apply_batch(merkleized).await.unwrap();
            db.commit().await.unwrap();

            // Mix in absent keys so some probes resolve to nothing.
            for _ in 0..100 {
                keys.push(Digest::random(&mut rng));
            }
            let refs: Vec<&Digest> = keys.iter().collect();
            let fused = db.get_many(&refs).await.unwrap();
            assert_eq!(fused.len(), keys.len());
            for (key, result) in keys.iter().zip(fused) {
                assert_eq!(result, db.get(key).await.unwrap());
            }

            db.destroy().await.unwrap();
        });
    }

    /// A merkleize future dropped mid-flight leaves the database fully usable.
    ///
    /// With a parallel strategy, merkleize's hashing job may keep running detached against its
    /// snapshot of committed Merkle state after the future is dropped. Later mutations must not
    /// observe it (they copy-on-write instead), and its discarded result must not corrupt
    /// anything. Runs on the tokio runtime for the same reason as
    /// [`test_get_many_fused_sharded_matches_get`].
    #[test_traced]
    fn test_merkleize_cancellation_leaves_db_usable() {
        let executor = commonware_runtime::tokio::Runner::default();
        executor.start(|context| async move {
            type ParTest = Db<
                mmr::Family,
                commonware_runtime::tokio::Context,
                Digest,
                Digest,
                Sha256,
                TwoCap,
                Rayon,
            >;
            let strategy = context.strategy(NZUsize!(2));
            let cfg = fixed_db_config_with_strategy::<TwoCap, Rayon>("cancel", &context, strategy);
            let mut db = ParTest::init(context.child("db"), cfg).await.unwrap();

            // Populate and commit so later snapshots carry committed nodes.
            let mut rng = TestRng::new(11);
            let mut keys = Vec::with_capacity(1024);
            let mut batch = db.new_batch();
            for _ in 0..1024 {
                let key = Digest::random(&mut rng);
                let value = Digest::random(&mut rng);
                keys.push((key, value));
                batch = batch.write(key, Some(value));
            }
            let merkleized = batch.merkleize(&db, None).await.unwrap();
            db.apply_batch(merkleized).await.unwrap();
            db.commit().await.unwrap();

            // Drop one merkleize after a single poll and race another against a timer, so the
            // future is abandoned at whatever stage it reached (possibly mid-hashing-job).
            for delay in [None, Some(Duration::from_millis(1))] {
                // Fork the abandoned batch off a two-deep unapplied chain: the hashing job
                // reaches the grandparent only through Weak references, so dropping the chain
                // below exercises a mid-job truncation (any resulting panic is caught and
                // discarded by `Strategy::spawn`).
                let grandparent = db
                    .new_batch()
                    .write(Digest::random(&mut rng), Some(Digest::random(&mut rng)))
                    .merkleize(&db, None)
                    .await
                    .unwrap();
                let parent = grandparent
                    .new_batch::<Sha256>()
                    .write(Digest::random(&mut rng), Some(Digest::random(&mut rng)))
                    .merkleize(&db, None)
                    .await
                    .unwrap();
                let mut abandoned = parent.new_batch::<Sha256>();
                for _ in 0..4200 {
                    abandoned =
                        abandoned.write(Digest::random(&mut rng), Some(Digest::random(&mut rng)));
                }
                let fut = abandoned.merkleize(&db, None);
                match delay {
                    None => {
                        let _ = fut.now_or_never();
                    }
                    Some(delay) => {
                        select! {
                            _ = fut => {},
                            _ = context.sleep(delay) => {},
                        }
                    }
                }
                drop(parent);
                drop(grandparent);

                // The database remains fully usable: mutate, merkleize, apply, and read back.
                let (key, _) = keys[0];
                let value = Digest::random(&mut rng);
                let merkleized = db
                    .new_batch()
                    .write(key, Some(value))
                    .merkleize(&db, None)
                    .await
                    .unwrap();
                db.apply_batch(merkleized).await.unwrap();
                db.commit().await.unwrap();
                assert_eq!(db.get(&key).await.unwrap(), Some(value));
                keys[0].1 = value;
                for (key, value) in &keys[1..] {
                    assert_eq!(db.get(key).await.unwrap(), Some(*value));
                }
            }

            db.destroy().await.unwrap();
        });
    }

    /// Create n random operations using the default seed (0). Some portion of
    /// the updates are deletes. create_test_ops(n) is a prefix of
    /// create_test_ops(n') for n < n'.
    pub(crate) fn create_test_ops(n: usize) -> Vec<Operation<mmr::Family, Digest, Digest>> {
        create_test_ops_seeded(n, 0)
    }

    /// Create n random operations using a specific seed.
    /// Use different seeds when you need non-overlapping keys in the same test.
    pub(crate) fn create_test_ops_seeded(
        n: usize,
        seed: u64,
    ) -> Vec<Operation<mmr::Family, Digest, Digest>> {
        let mut rng = TestRng::new(seed);
        let mut prev_key = Digest::random(&mut rng);
        let mut ops = Vec::new();
        for i in 0..n {
            let key = Digest::random(&mut rng);
            if i % 10 == 0 && i > 0 {
                ops.push(Operation::Delete(prev_key));
            } else {
                let value = Digest::random(&mut rng);
                ops.push(Operation::Update(Update(key, value)));
                prev_key = key;
            }
        }
        ops
    }

    /// Applies the given operations to the database.
    pub(crate) async fn apply_ops(
        db: &mut AnyTest,
        ops: Vec<Operation<mmr::Family, Digest, Digest>>,
    ) {
        let mut batch = db.new_batch();
        for op in ops {
            match op {
                Operation::Update(Update(key, value)) => {
                    batch = batch.write(key, Some(value));
                }
                Operation::Delete(key) => {
                    batch = batch.write(key, None);
                }
                Operation::CommitFloor(_, _) => {
                    panic!("CommitFloor not supported in apply_ops");
                }
            }
        }
        let merkleized = batch.merkleize(db, None).await.unwrap();
        db.apply_batch(merkleized).await.unwrap();
    }

    /// Helper: commit a batch of key-value writes and return the applied range (generic).
    async fn commit_writes_generic<F: Family>(
        db: &mut AnyTestGeneric<F>,
        writes: impl IntoIterator<Item = (Digest, Option<Digest>)>,
        metadata: Option<Digest>,
    ) -> std::ops::Range<GenericLocation<F>> {
        let mut batch = db.new_batch();
        for (k, v) in writes {
            batch = batch.write(k, v);
        }
        let merkleized = batch.merkleize(db, metadata).await.unwrap();
        let range = db.apply_batch(merkleized).await.unwrap();
        db.commit().await.unwrap();
        range
    }

    fn key(i: u64) -> Digest {
        Sha256::hash(&i.to_be_bytes())
    }

    fn val(i: u64) -> Digest {
        Sha256::hash(&(i + 10000).to_be_bytes())
    }

    /// The init-time `(location -> key)` cache only memoizes log reads, so rebuilding the snapshot
    /// with the cache disabled (`init_cache_size = None`) or enabled must produce the identical root.
    #[test_traced("WARN")]
    fn test_unordered_fixed_init_cache_equivalence() {
        deterministic::Runner::default().start(|context| async move {
            // Populate a database with churny operations (repeated updates and deletes drive the
            // collision resolution that the cache accelerates), then commit and drop it.
            let cfg = fixed_db_config::<TwoCap>("cache_equiv", &context);
            let mut db = AnyTest::init(context.child("populate"), cfg).await.unwrap();
            apply_ops(&mut db, create_test_ops(10_000)).await;
            db.commit().await.unwrap();
            db.sync().await.unwrap();
            let root = db.root();
            drop(db);

            // Reopen with the cache disabled and with a large cache; both rebuild the snapshot by
            // replaying the same immutable log, so both roots must equal the pre-drop root.
            for cache_size in [None, Some(NZUsize!(1 << 20))] {
                let mut cfg = fixed_db_config::<TwoCap>("cache_equiv", &context);
                cfg.init_cache_size = cache_size;
                let ctx = context
                    .child("reopen")
                    .with_attribute("cache", cache_size.map_or(0, NonZeroUsize::get));
                let db = AnyTest::init(ctx, cfg).await.unwrap();
                assert_eq!(
                    db.root(),
                    root,
                    "root mismatch at cache_size={cache_size:?}"
                );
                drop(db);
            }
        });
    }

    #[test_traced("INFO")]
    fn test_any_unordered_fixed_metrics() {
        deterministic::Runner::default().start(|ctx| async move {
            let mut db = open_db_generic::<mmr::Family>(ctx.child("db")).await;
            let k = key(1);
            let v = val(1);
            let batch = db
                .new_batch()
                .write(k, Some(v))
                .merkleize(&db, None)
                .await
                .unwrap();
            db.apply_batch(batch).await.unwrap();
            assert_eq!(db.get(&k).await.unwrap(), Some(v));
            assert_eq!(db.get_many(&[&k]).await.unwrap(), vec![Some(v)]);
            db.commit().await.unwrap();
            db.sync().await.unwrap();
            db.prune(Location::new(0)).await.unwrap();

            let metrics = ctx.encode();
            for expected in [
                "db_size 4",
                "db_pruning_boundary 0",
                "db_retained 4",
                "db_inactivity_floor 2",
                "db_last_commit 3",
                "db_get_calls_total 1",
                "db_get_many_calls_total 1",
                "db_lookups_requested_total 2",
                "db_apply_batch_calls_total 1",
                "db_operations_applied_total 3",
                "db_commit_calls_total 1",
                "db_sync_calls_total 1",
                "db_prune_calls_total 1",
                "db_get_duration_count 1",
                "db_get_many_duration_count 1",
                "db_apply_batch_duration_count 1",
                "db_commit_duration_count 1",
                "db_sync_duration_count 1",
                "db_prune_duration_count 1",
            ] {
                assert!(metrics.contains(expected), "missing {expected}\n{metrics}");
            }
        });
    }

    /// Reads on a batch must not perturb `merkleize`: the root must be byte-identical to a
    /// write-only batch's `merkleize`, across updates/deletes/creates, both with the batch
    /// rooted directly at the DB (D=0) and through pending ancestors (D=1, D=2).
    #[test_traced("WARN")]
    fn test_unordered_fixed_read_merkleize_parity() {
        type ParentChain = Vec<
            std::sync::Arc<
                crate::qmdb::any::batch::MerkleizedBatch<
                    mmr::Family,
                    Digest,
                    crate::qmdb::any::operation::update::Unordered<Digest, FixedEncoding<Digest>>,
                    Sequential,
                >,
            >,
        >;

        deterministic::Runner::default().start(|ctx| async move {
            let mut db = create_test_db(ctx.child("db")).await;

            // Seed 2000 keys and commit so they live in the committed snapshot.
            let mut seed = db.new_batch();
            for i in 0..2000u64 {
                seed = seed.write(key(i), Some(val(i)));
            }
            let seed = seed.merkleize(&db, None).await.unwrap();
            db.apply_batch(seed).await.unwrap();
            db.commit().await.unwrap();

            // Build a mixed mutation set: updates of existing keys, deletes of existing keys,
            // and creates of fresh keys. `make` re-derives the set from a seed so both paths and
            // both depths see identical mutations.
            let make = |salt: u64| -> Vec<(Digest, Option<Digest>)> {
                let mut rng = TestRng::new(salt);
                let mut out = Vec::new();
                for _ in 0..600 {
                    let r = rng.next_u32() % 100;
                    if r < 60 {
                        out.push((key(rng.next_u64() % 2000), Some(val(rng.next_u64()))));
                    } else if r < 80 {
                        out.push((key(rng.next_u64() % 2000), None));
                    } else {
                        out.push((key(2000 + rng.next_u64() % 2000), Some(val(rng.next_u64()))));
                    }
                }
                // Dedup last-write-wins.
                let mut m: HashMap<Digest, Option<Digest>> = HashMap::new();
                for (k, v) in out {
                    m.insert(k, v);
                }
                m.into_iter().collect()
            };

            // D=0: batch rooted directly at the DB. D=N: through N pending ancestors.
            for depth in [0u64, 1, 2] {
                let mut chain: ParentChain = Vec::new();
                for d in 0..depth {
                    let mut p = chain
                        .last()
                        .map_or_else(|| db.new_batch(), |l| l.new_batch::<Sha256>());
                    for (k, v) in make(900 + d) {
                        p = p.write(k, v);
                    }
                    chain.push(p.merkleize(&db, None).await.unwrap());
                }

                let muts = make(depth + 1);
                let new_batch = || {
                    chain
                        .last()
                        .map_or_else(|| db.new_batch(), |p| p.new_batch::<Sha256>())
                };

                // Normal path.
                let mut nb = new_batch();
                for (k, v) in &muts {
                    nb = nb.write(*k, *v);
                }
                let normal_root = nb.merkleize(&db, None).await.unwrap().root();

                // Read-then-write on one batch. Values and root must match the write-only path.
                let keys: Vec<&Digest> = muts.iter().map(|(k, _)| k).collect();
                let mut fb = new_batch();
                // Duplicate keys in one read resolve identically per slot.
                let dup_values = fb.get_many(&[keys[0], keys[0]], &db).await.unwrap();
                assert_eq!(dup_values[0], dup_values[1]);
                // Keys read but never written must not affect the root.
                let unwritten: Vec<Digest> = (0..40u64)
                    .map(|i| key(i * 50))
                    .chain((0..5).map(|i| key(8000 + i)))
                    .collect();
                let unwritten_refs: Vec<&Digest> = unwritten.iter().collect();
                fb.get_many(&unwritten_refs, &db).await.unwrap();
                let values = fb.get_many(&keys, &db).await.unwrap();
                let plain = new_batch().get_many(&keys, &db).await.unwrap();
                assert_eq!(values, plain, "value mismatch at depth={depth}");
                for (k, v) in &muts {
                    fb = fb.write(*k, *v);
                }
                let fused_root = fb.merkleize(&db, None).await.unwrap().root();
                assert_eq!(normal_root, fused_root, "root mismatch at depth={depth}");

                // Reads after writes: written keys are answered by the pending mutations and the
                // root must still match.
                let half = muts.len() / 2;
                let mut mb = new_batch();
                for (k, v) in muts.iter().take(half) {
                    mb = mb.write(*k, *v);
                }
                let values = mb.get_many(&keys, &db).await.unwrap();
                for (i, (_, v)) in muts.iter().enumerate().take(half) {
                    assert_eq!(values[i], *v, "pending write not visible at depth={depth}");
                }
                assert_eq!(
                    values[half..],
                    plain[half..],
                    "unwritten value mismatch at depth={depth}"
                );
                for (k, v) in muts.iter().skip(half) {
                    mb = mb.write(*k, *v);
                }
                let mixed_root = mb.merkleize(&db, None).await.unwrap().root();
                assert_eq!(
                    normal_root, mixed_root,
                    "mixed root mismatch at depth={depth}"
                );

                // Multiple disjoint reads and single-key gets must not affect the root.
                let mut gb = new_batch();
                gb.get_many(&keys[..half], &db).await.unwrap();
                for key in &keys[half..] {
                    gb.get(key, &db).await.unwrap();
                }
                for (k, v) in &muts {
                    gb = gb.write(*k, *v);
                }
                let merged_root = gb.merkleize(&db, None).await.unwrap().root();
                assert_eq!(
                    normal_root, merged_root,
                    "merged root mismatch at depth={depth}"
                );
            }
        });
    }

    // -- Generic inner functions for parameterized batch tests --

    async fn batch_empty_inner<F: Family>(context: deterministic::Context) {
        let mut db = open_db_generic::<F>(context.child("db")).await;
        let root_before = db.root();

        let merkleized = db.new_batch().merkleize(&db, None).await.unwrap();
        db.apply_batch(merkleized).await.unwrap();
        assert_ne!(db.root(), root_before);

        // DB should still be functional.
        commit_writes_generic(&mut db, [(key(0), Some(val(0)))], None).await;
        assert_eq!(db.get(&key(0)).await.unwrap(), Some(val(0)));

        db.destroy().await.unwrap();
    }

    async fn batch_metadata_inner<F: Family>(context: deterministic::Context) {
        let mut db = open_db_generic::<F>(context.child("db")).await;
        let metadata = val(42);

        commit_writes_generic(&mut db, [(key(0), Some(val(0)))], Some(metadata)).await;
        assert_eq!(db.get_metadata().await.unwrap(), Some(metadata));

        let merkleized = db.new_batch().merkleize(&db, None).await.unwrap();
        db.apply_batch(merkleized).await.unwrap();
        assert_eq!(db.get_metadata().await.unwrap(), None);

        db.destroy().await.unwrap();
    }

    async fn batch_get_read_through_inner<F: Family>(context: deterministic::Context) {
        let mut db = open_db_generic::<F>(context.child("db")).await;

        let ka = key(0);
        let va = val(0);
        commit_writes_generic(&mut db, [(ka, Some(va))], None).await;

        let kb = key(1);
        let vb = val(1);
        let kc = key(2);

        let mut batch = db.new_batch();
        assert_eq!(batch.get(&ka, &db).await.unwrap(), Some(va));

        batch = batch.write(kb, Some(vb));
        assert_eq!(batch.get(&kb, &db).await.unwrap(), Some(vb));
        assert_eq!(batch.get(&kc, &db).await.unwrap(), None);

        let va2 = val(100);
        batch = batch.write(ka, Some(va2));
        assert_eq!(batch.get(&ka, &db).await.unwrap(), Some(va2));

        batch = batch.write(ka, None);
        assert_eq!(batch.get(&ka, &db).await.unwrap(), None);

        db.destroy().await.unwrap();
    }

    async fn batch_get_on_merkleized_inner<F: Family>(context: deterministic::Context) {
        let mut db = open_db_generic::<F>(context.child("db")).await;

        let ka = key(0);
        let kb = key(1);
        let kc = key(2);
        let kd = key(3);

        commit_writes_generic(&mut db, [(ka, Some(val(0))), (kb, Some(val(1)))], None).await;

        let va2 = val(100);
        let vc = val(2);
        let merkleized = db
            .new_batch()
            .write(ka, Some(va2))
            .write(kb, None)
            .write(kc, Some(vc))
            .merkleize(&db, None)
            .await
            .unwrap();

        assert_eq!(merkleized.get(&ka, &db).await.unwrap(), Some(va2));
        assert_eq!(merkleized.get(&kb, &db).await.unwrap(), None);
        assert_eq!(merkleized.get(&kc, &db).await.unwrap(), Some(vc));
        assert_eq!(merkleized.get(&kd, &db).await.unwrap(), None);

        db.destroy().await.unwrap();
    }

    async fn batch_stacked_get_inner<F: Family>(context: deterministic::Context) {
        let db = open_db_generic::<F>(context.child("db")).await;

        let ka = key(0);
        let kb = key(1);

        let merkleized = db
            .new_batch()
            .write(ka, Some(val(0)))
            .merkleize(&db, None)
            .await
            .unwrap();

        let mut child = merkleized.new_batch();
        assert_eq!(child.get(&ka, &db).await.unwrap(), Some(val(0)));

        child = child.write(ka, Some(val(100)));
        assert_eq!(child.get(&ka, &db).await.unwrap(), Some(val(100)));

        child = child.write(kb, Some(val(1)));
        assert_eq!(child.get(&kb, &db).await.unwrap(), Some(val(1)));

        child = child.write(ka, None);
        assert_eq!(child.get(&ka, &db).await.unwrap(), None);

        drop(child);
        drop(merkleized);
        db.destroy().await.unwrap();
    }

    async fn batch_stacked_delete_recreate_inner<F: Family>(context: deterministic::Context) {
        let mut db = open_db_generic::<F>(context.child("db")).await;
        let ka = key(0);

        commit_writes_generic(&mut db, [(ka, Some(val(0)))], None).await;

        let parent_m = db
            .new_batch()
            .write(ka, None)
            .merkleize(&db, None)
            .await
            .unwrap();
        assert_eq!(parent_m.get(&ka, &db).await.unwrap(), None);

        let child_m = parent_m
            .new_batch()
            .write(ka, Some(val(200)))
            .merkleize(&db, None)
            .await
            .unwrap();
        assert_eq!(child_m.get(&ka, &db).await.unwrap(), Some(val(200)));

        db.apply_batch(child_m).await.unwrap();
        assert_eq!(db.get(&ka).await.unwrap(), Some(val(200)));

        db.destroy().await.unwrap();
    }

    async fn batch_apply_returns_range_inner<F: Family>(context: deterministic::Context) {
        let mut db = open_db_generic::<F>(context.child("db")).await;

        let writes: Vec<_> = (0..5).map(|i| (key(i), Some(val(i)))).collect();
        let range1 = commit_writes_generic(&mut db, writes, None).await;

        assert_eq!(range1.start, GenericLocation::<F>::new(1));
        assert!(range1.end.saturating_sub(*range1.start) >= 6);

        let writes: Vec<_> = (5..10).map(|i| (key(i), Some(val(i)))).collect();
        let range2 = commit_writes_generic(&mut db, writes, None).await;
        assert_eq!(range2.start, range1.end);

        db.destroy().await.unwrap();
    }

    async fn batch_speculative_root_inner<F: Family>(context: deterministic::Context) {
        let mut db = open_db_generic::<F>(context.child("db")).await;

        let mut batch = db.new_batch();
        for i in 0..10 {
            batch = batch.write(key(i), Some(val(i)));
        }
        let merkleized = batch.merkleize(&db, None).await.unwrap();
        let speculative_root = merkleized.root();

        db.apply_batch(merkleized).await.unwrap();
        assert_eq!(db.root(), speculative_root);

        db.destroy().await.unwrap();
    }

    async fn log_replay_inner<F: Family>(context: deterministic::Context) {
        let db_context = context.child("db");
        let mut db = open_db_generic::<F>(db_context.child("db")).await;

        // Update the same key many times within a single batch.
        const UPDATES: u64 = 100;
        let k = Sha256::hash(&UPDATES.to_be_bytes());
        let mut batch = db.new_batch();
        for i in 0u64..UPDATES {
            let v = Sha256::hash(&(i * 1000).to_be_bytes());
            batch = batch.write(k, Some(v));
        }
        let merkleized = batch.merkleize(&db, None).await.unwrap();
        db.apply_batch(merkleized).await.unwrap();
        db.commit().await.unwrap();
        let root = db.root();

        // Simulate a failed commit and test that the log replay doesn't leave behind old data.
        drop(db);
        let db: AnyTestGeneric<F> = open_db_generic::<F>(db_context.child("reopened")).await;
        let iter = db.snapshot.get(&k);
        assert_eq!(iter.cloned().collect::<Vec<_>>().len(), 1);
        assert_eq!(db.root(), root);

        db.destroy().await.unwrap();
    }

    // -- MMR test wrappers --

    #[test_traced("INFO")]
    fn test_unordered_fixed_batch_empty() {
        let executor = deterministic::Runner::default();
        executor.start(batch_empty_inner::<mmr::Family>);
    }

    #[test_traced("INFO")]
    fn test_unordered_fixed_batch_metadata() {
        let executor = deterministic::Runner::default();
        executor.start(batch_metadata_inner::<mmr::Family>);
    }

    #[test_traced("INFO")]
    fn test_unordered_fixed_batch_get_read_through() {
        let executor = deterministic::Runner::default();
        executor.start(batch_get_read_through_inner::<mmr::Family>);
    }

    #[test_traced("INFO")]
    fn test_unordered_fixed_batch_get_on_merkleized() {
        let executor = deterministic::Runner::default();
        executor.start(batch_get_on_merkleized_inner::<mmr::Family>);
    }

    #[test_traced("INFO")]
    fn test_unordered_fixed_batch_stacked_get() {
        let executor = deterministic::Runner::default();
        executor.start(batch_stacked_get_inner::<mmr::Family>);
    }

    #[test_traced("INFO")]
    fn test_unordered_fixed_batch_stacked_delete_recreate() {
        let executor = deterministic::Runner::default();
        executor.start(batch_stacked_delete_recreate_inner::<mmr::Family>);
    }

    #[test_traced("INFO")]
    fn test_unordered_fixed_batch_apply_returns_range() {
        let executor = deterministic::Runner::default();
        executor.start(batch_apply_returns_range_inner::<mmr::Family>);
    }

    #[test_traced("INFO")]
    fn test_unordered_fixed_batch_speculative_root() {
        let executor = deterministic::Runner::default();
        executor.start(batch_speculative_root_inner::<mmr::Family>);
    }

    #[test_traced("WARN")]
    fn test_any_fixed_db_log_replay() {
        let executor = deterministic::Runner::default();
        executor.start(log_replay_inner::<mmr::Family>);
    }

    // -- MMB test wrappers --

    #[test_traced("INFO")]
    fn test_unordered_fixed_batch_empty_mmb() {
        let executor = deterministic::Runner::default();
        executor.start(batch_empty_inner::<crate::merkle::mmb::Family>);
    }

    #[test_traced("INFO")]
    fn test_unordered_fixed_batch_metadata_mmb() {
        let executor = deterministic::Runner::default();
        executor.start(batch_metadata_inner::<crate::merkle::mmb::Family>);
    }

    #[test_traced("INFO")]
    fn test_unordered_fixed_batch_get_read_through_mmb() {
        let executor = deterministic::Runner::default();
        executor.start(batch_get_read_through_inner::<crate::merkle::mmb::Family>);
    }

    #[test_traced("INFO")]
    fn test_unordered_fixed_batch_get_on_merkleized_mmb() {
        let executor = deterministic::Runner::default();
        executor.start(batch_get_on_merkleized_inner::<crate::merkle::mmb::Family>);
    }

    #[test_traced("INFO")]
    fn test_unordered_fixed_batch_stacked_get_mmb() {
        let executor = deterministic::Runner::default();
        executor.start(batch_stacked_get_inner::<crate::merkle::mmb::Family>);
    }

    #[test_traced("INFO")]
    fn test_unordered_fixed_batch_stacked_delete_recreate_mmb() {
        let executor = deterministic::Runner::default();
        executor.start(batch_stacked_delete_recreate_inner::<crate::merkle::mmb::Family>);
    }

    #[test_traced("INFO")]
    fn test_unordered_fixed_batch_apply_returns_range_mmb() {
        let executor = deterministic::Runner::default();
        executor.start(batch_apply_returns_range_inner::<crate::merkle::mmb::Family>);
    }

    #[test_traced("INFO")]
    fn test_unordered_fixed_batch_speculative_root_mmb() {
        let executor = deterministic::Runner::default();
        executor.start(batch_speculative_root_inner::<crate::merkle::mmb::Family>);
    }

    #[test_traced("WARN")]
    fn test_any_fixed_db_log_replay_mmb() {
        let executor = deterministic::Runner::default();
        executor.start(log_replay_inner::<crate::merkle::mmb::Family>);
    }

    // -- MMR-only tests (use verify_proof / Position which are MMR-specific) --

    #[test]
    fn test_any_fixed_db_historical_proof_basic() {
        let executor = deterministic::Runner::default();
        executor.start(|context| async move {
            let mut db = create_test_db(context.child("storage")).await;
            let ops = create_test_ops(20);
            apply_ops(&mut db, ops.clone()).await;
            let root_hash = db.root();
            let original_op_count = db.bounds().end;

            // Historical proof should match "regular" proof when historical size == current database size
            let max_ops = NZU64!(10);
            let (historical_proof, historical_ops) = db
                .historical_proof(original_op_count, Location::new(6), max_ops)
                .await
                .unwrap();
            let (regular_proof, regular_ops) = db.proof(Location::new(6), max_ops).await.unwrap();

            assert_eq!(historical_proof.leaves, regular_proof.leaves);
            assert_eq!(historical_proof.digests, regular_proof.digests);
            assert_eq!(historical_ops, regular_ops);
            assert!(verify_proof::<Sha256, _, _>(
                &historical_proof,
                Location::new(6),
                &historical_ops,
                &root_hash
            ));

            // Add more operations to the database
            // (use different seed to avoid key collisions)
            let more_ops = create_test_ops_seeded(5, 1);
            apply_ops(&mut db, more_ops.clone()).await;

            // Historical proof should remain the same even though database has grown
            let (historical_proof, historical_ops) = db
                .historical_proof(original_op_count, Location::new(6), NZU64!(10))
                .await
                .unwrap();
            assert_eq!(historical_proof.leaves, original_op_count);
            assert_eq!(historical_proof.leaves, regular_proof.leaves);
            assert_eq!(historical_ops.len(), 10);
            assert_eq!(historical_proof.digests, regular_proof.digests);
            assert_eq!(historical_ops, regular_ops);
            assert!(verify_proof::<Sha256, _, _>(
                &historical_proof,
                Location::new(6),
                &historical_ops,
                &root_hash
            ));

            // Try to get historical proof with op_count > number of operations and confirm it
            // returns RangeOutOfBounds error.
            assert!(matches!(
                db.historical_proof(db.bounds().end + 1, Location::new(6), NZU64!(10))
                    .await,
                Err(Error::Merkle(crate::mmr::Error::RangeOutOfBounds(_)))
            ));

            db.destroy().await.unwrap();
        });
    }

    #[test]
    fn test_any_fixed_db_historical_proof_edge_cases() {
        let executor = deterministic::Runner::default();
        executor.start(|context| async move {
            let mut db = create_test_db(context.child("first")).await;
            // Apply ops in multiple batches; each apply_ops ends in a commit, so the size
            // after each batch is a commit-boundary historical size.
            let mut commit_boundary_sizes: Vec<Location> = Vec::new();
            for _ in 0..5 {
                apply_ops(&mut db, create_test_ops(10)).await;
                commit_boundary_sizes.push(db.bounds().end);
            }

            let root = db.root();
            let full_size = db.bounds().end;
            assert_eq!(full_size, *commit_boundary_sizes.last().unwrap());

            // Verify a single-op proof at the full commit size.
            let (proof, proof_ops) = db.proof(Location::new(1), NZU64!(1)).await.unwrap();
            assert_eq!(proof_ops.len(), 1);
            assert!(verify_proof::<Sha256, _, _>(
                &proof,
                Location::new(1),
                &proof_ops,
                &root
            ));

            // historical_proof at full size should match proof.
            let (hp, hp_ops) = db
                .historical_proof(full_size, Location::new(1), NZU64!(1))
                .await
                .unwrap();
            assert_eq!(hp.digests, proof.digests);
            assert_eq!(hp_ops, proof_ops);

            // Test requesting more operations than available in historical position. Use
            // the commit-boundary size after the first batch.
            let first_batch_size = commit_boundary_sizes[0];
            let (_proof, limited_ops) = db
                .historical_proof(first_batch_size, Location::new(6), NZU64!(1000))
                .await
                .unwrap();
            assert_eq!(limited_ops.len() as u64, *first_batch_size - 6);

            // Test proof at minimum historical position (just the initial commit).
            let (min_proof, min_ops) = db
                .historical_proof(Location::new(1), Location::new(0), NZU64!(3))
                .await
                .unwrap();
            assert_eq!(min_proof.leaves, Location::new(1));
            assert_eq!(min_ops.len(), 1);

            // Non-commit-boundary historical sizes are rejected.
            let result = db
                .historical_proof(Location::new(5), Location::new(1), NZU64!(3))
                .await;
            assert!(
                matches!(result, Err(crate::qmdb::Error::HistoricalFloorPruned(loc)) if loc == Location::new(5)),
                "expected HistoricalFloorPruned(5), got {result:?}"
            );

            db.destroy().await.unwrap();
        });
    }

    #[test]
    fn test_any_fixed_db_historical_proof_different_historical_sizes() {
        let executor = deterministic::Runner::default();
        executor.start(|context| async move {
            let ops = create_test_ops(100);
            let start_loc = Location::new(2);
            let max_ops = NZU64!(10);

            // Build checkpoints only at commit points and record reference proofs/roots there.
            let mut db = create_test_db(context.child("main")).await;
            let mut offset = 0usize;
            let mut checkpoints = Vec::new();
            for chunk in [20usize, 15, 25, 30, 10] {
                apply_ops(&mut db, ops[offset..offset + chunk].to_vec()).await;
                offset += chunk;

                let end_loc = db.bounds().end;
                let root = db.root();
                let (proof, proof_ops) = db.proof(start_loc, max_ops).await.unwrap();
                checkpoints.push((end_loc, root, proof, proof_ops));
            }

            // Grow state past the checkpoints with an empty batch and verify all
            // historical proofs from that later state.
            let merkleized = db.new_batch().merkleize(&db, None).await.unwrap();
            db.apply_batch(merkleized).await.unwrap();
            for (historical_size, root, reference_proof, reference_ops) in checkpoints {
                let (historical_proof, historical_ops) = db
                    .historical_proof(historical_size, start_loc, max_ops)
                    .await
                    .unwrap();
                assert_eq!(historical_proof.leaves, reference_proof.leaves);
                assert_eq!(historical_proof.digests, reference_proof.digests);
                assert_eq!(historical_ops, reference_ops);
                assert!(verify_proof::<Sha256, _, _>(
                    &historical_proof,
                    start_loc,
                    &historical_ops,
                    &root
                ));
            }

            // Verify the current full-size proof against the current root as a final sanity check.
            let full_root = db.root();
            let (full_proof, full_ops) = db.proof(start_loc, max_ops).await.unwrap();
            assert!(verify_proof::<Sha256, _, _>(
                &full_proof,
                start_loc,
                &full_ops,
                &full_root
            ));

            db.destroy().await.unwrap();
        });
    }

    fn is_send<T: Send>(_: T) {}

    #[allow(dead_code)]
    fn assert_non_trait_futures_are_send(db: &AnyTest, key: Digest, value: Digest) {
        let reader = db.new_batch();
        is_send(reader.get_many(&[&key], db));
        let batch = db.new_batch().write(key, Some(value));
        is_send(batch.merkleize(db, None));
        is_send(db.get_with_loc(&key));
    }

    // FromSyncTestable implementation for from_sync_result tests
    mod from_sync_testable {
        use super::*;
        use crate::{
            merkle::mmr::{self, full::Mmr},
            qmdb::any::sync::tests::FromSyncTestable,
        };
        use futures::future::join_all;

        type TestMmr = Mmr<deterministic::Context, Digest, Sequential>;

        impl FromSyncTestable for AnyTest {
            type Merkle = TestMmr;

            fn into_log_components(self) -> (Self::Merkle, Self::Journal) {
                (self.log.merkle, self.log.journal)
            }

            async fn pinned_nodes_at(&self, loc: Location) -> Vec<Digest> {
                join_all(mmr::Family::nodes_to_pin(loc).map(|p| self.log.merkle.get_node(p)))
                    .await
                    .into_iter()
                    .map(|n| n.unwrap().unwrap())
                    .collect()
            }
        }
    }
}