//! An MMR backed by a fixed-item-length journal.
//!
//! A [crate::journal] is used to store all unpruned MMR nodes, and a [crate::metadata] store is
//! used to preserve digests required for root and proof generation that would have otherwise been
//! pruned.
//!
//! This module is a thin wrapper around the generic `Merkle` type, specialized for the
//! MMR [Family]. It re-exports [Config], [SyncConfig], and the append-only
//! [`UnmerkleizedBatch`] wrapper from `merkle::full`. Async proof methods (`proof`,
//! `range_proof`, `historical_proof`, `historical_range_proof`) and the `Storage<F>` impl are
//! provided by the generic `Merkle` in `merkle::full`.

/// Configuration for a journal-backed MMR.
pub use crate::merkle::full::Config;
pub use crate::merkle::full::UnmerkleizedBatch;
use crate::merkle::mmr::Family;

/// Configuration for initializing a full MMR for synchronization.
pub type SyncConfig<D, S> = crate::merkle::full::SyncConfig<Family, D, S>;

/// A MMR backed by a fixed-item-length journal.
pub type Mmr<E, D, S> = crate::merkle::full::Merkle<Family, E, D, S>;

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{
        merkle::{Bagging::ForwardFold, Family as _, conformance::build_test_mmr},
        mmr::{Error, Location, Position, StandardHasher as Standard, mem},
    };
    use commonware_cryptography::{Hasher, Sha256, sha256::Digest};
    use commonware_macros::test_traced;
    use commonware_parallel::Sequential;
    use commonware_runtime::{Runner, Supervisor as _, buffer::paged::CacheRef, deterministic};
    use commonware_utils::{NZU16, NZU64, NZUsize, non_empty_range};
    use rstest::rstest;
    use std::num::{NonZeroU16, NonZeroUsize};

    fn test_digest(v: usize) -> Digest {
        Sha256::hash(&[&v.to_be_bytes()])
    }

    const PAGE_SIZE: NonZeroU16 = NZU16!(111);
    const PAGE_CACHE_SIZE: NonZeroUsize = NZUsize!(5);

    fn test_config() -> Config<Sequential> {
        Config {
            journal_partition: "journal-partition".into(),
            metadata_partition: "metadata-partition".into(),
            items_per_blob: NZU64!(7),
            write_buffer: NZUsize!(1024),
            replay_buffer: NZUsize!(1024),
            strategy: Sequential,
            page_cache: CacheRef::new(PAGE_SIZE, PAGE_CACHE_SIZE),
        }
    }

    /// Test that the full MMR produces the same root as the in-memory reference.
    #[test]
    fn test_full_mmr_batched_root() {
        let executor = deterministic::Runner::default();
        executor.start(|context| async move {
            const NUM_ELEMENTS: u64 = 199;
            let hasher: Standard<Sha256> = Standard::new(ForwardFold);
            let test_mmr = mem::Mmr::new();
            let test_mmr = build_test_mmr(&hasher, test_mmr, NUM_ELEMENTS);
            let expected_root = test_mmr.root(&hasher, 0).unwrap();

            let mmr = Mmr::init(
                context.child("storage"),
                &Standard::<Sha256>::new(ForwardFold),
                test_config(),
            )
            .await
            .unwrap();

            let mut batch = mmr.new_batch();
            for i in 0u64..NUM_ELEMENTS {
                let element = hasher.digest(&i.to_be_bytes());
                batch = batch.add(&hasher, &element);
            }
            let batch = batch.merkleize(mmr.mem(), &hasher);
            let mmr = mmr.apply_batch(&batch).unwrap();
            assert_eq!(mmr.root(&hasher, 0).unwrap(), expected_root);

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

    #[test_traced]
    fn test_full_mmr_pop() {
        let executor = deterministic::Runner::default();
        executor.start(|context| async move {
            const NUM_ELEMENTS: u64 = 200;

            let hasher: Standard<Sha256> = Standard::new(ForwardFold);
            let cfg = test_config();
            let mut mmr = Mmr::init(context.child("init"), &hasher, cfg)
                .await
                .unwrap();

            let mut c_hasher = Sha256::default();
            let mut batch = mmr.new_batch();
            for i in 0u64..NUM_ELEMENTS {
                c_hasher.update(&i.to_be_bytes());
                let (next_hasher, element) = c_hasher.finalize();
                c_hasher = next_hasher;
                batch = batch.add(&hasher, &element);
            }
            let batch = batch.merkleize(mmr.mem(), &hasher);
            mmr = mmr.apply_batch(&batch).unwrap();

            // Sync and reopen one leaf earlier until empty, confirming the root each time.
            for i in (0..NUM_ELEMENTS).rev() {
                let cap = mmr.leaves() - 1;
                _ = mmr.sync().await.unwrap();
                mmr = Mmr::init_at_most(context.child("cap"), &hasher, test_config(), cap)
                    .await
                    .unwrap();
                let root = mmr.root(&hasher, 0).unwrap();
                let mut reference_mmr = mem::Mmr::new();
                let batch = {
                    let mut batch = reference_mmr.new_batch();
                    for j in 0..i {
                        c_hasher.update(&j.to_be_bytes());
                        let (next_hasher, element) = c_hasher.finalize();
                        c_hasher = next_hasher;
                        batch = batch.add(&hasher, &element);
                    }
                    batch.merkleize(&reference_mmr, &hasher)
                };
                reference_mmr.apply_batch(&batch).unwrap();
                assert_eq!(
                    root,
                    reference_mmr.root(&hasher, 0).unwrap(),
                    "root mismatch after bounded initialization at {i}"
                );
            }
            let cap = mmr.leaves();
            _ = mmr.sync().await.unwrap();
            mmr = Mmr::init_at_most(context.child("cap"), &hasher, test_config(), cap)
                .await
                .unwrap();
            drop(mmr);
            let mut mmr = Mmr::init(context.child("reopen1"), &hasher, test_config())
                .await
                .unwrap();

            // Repeat the test though sync part of the way to tip to test crossing the boundary from
            // cached to uncached leaves, and reduce the initialization bound by 2 at a time instead
            // of just 1.
            {
                let mut batch = mmr.new_batch();
                for i in 0u64..NUM_ELEMENTS {
                    c_hasher.update(&i.to_be_bytes());
                    let (next_hasher, element) = c_hasher.finalize();
                    c_hasher = next_hasher;
                    batch = batch.add(&hasher, &element);
                    if i == 101 {
                        // We can't sync mid-batch, so apply the first part,
                        // sync, then start a new batch for the rest.
                        break;
                    }
                }
                let batch = batch.merkleize(mmr.mem(), &hasher);
                mmr = mmr.apply_batch(&batch).unwrap();
                mmr = mmr.sync().await.unwrap();
                let mut batch = mmr.new_batch();
                for i in 102u64..NUM_ELEMENTS {
                    c_hasher.update(&i.to_be_bytes());
                    let (next_hasher, element) = c_hasher.finalize();
                    c_hasher = next_hasher;
                    batch = batch.add(&hasher, &element);
                }
                let batch = batch.merkleize(mmr.mem(), &hasher);
                mmr = mmr.apply_batch(&batch).unwrap();
            }

            for i in (0..NUM_ELEMENTS - 1).rev().step_by(2) {
                let cap = mmr.leaves() - 2;
                _ = mmr.sync().await.unwrap();
                mmr = Mmr::init_at_most(context.child("cap"), &hasher, test_config(), cap)
                    .await
                    .unwrap();
                let root = mmr.root(&hasher, 0).unwrap();
                let reference_mmr = mem::Mmr::new();
                let reference_mmr = build_test_mmr(&hasher, reference_mmr, i);
                assert_eq!(
                    root,
                    reference_mmr.root(&hasher, 0).unwrap(),
                    "root mismatch at position {i:?}"
                );
            }
            // Persist the empty state so re-initialization below starts from it.
            mmr = mmr.sync().await.unwrap();
            drop(mmr);
            let mut mmr = Mmr::init(context.child("reopen2"), &hasher, test_config())
                .await
                .unwrap();

            // Repeat one more time only after pruning the MMR first.
            {
                let mut batch = mmr.new_batch();
                for i in 0u64..102 {
                    c_hasher.update(&i.to_be_bytes());
                    let (next_hasher, element) = c_hasher.finalize();
                    c_hasher = next_hasher;
                    batch = batch.add(&hasher, &element);
                }
                let batch = batch.merkleize(mmr.mem(), &hasher);
                mmr = mmr.apply_batch(&batch).unwrap();
                mmr = mmr.sync().await.unwrap();
                let mut batch = mmr.new_batch();
                for i in 102u64..NUM_ELEMENTS {
                    c_hasher.update(&i.to_be_bytes());
                    let (next_hasher, element) = c_hasher.finalize();
                    c_hasher = next_hasher;
                    batch = batch.add(&hasher, &element);
                }
                let batch = batch.merkleize(mmr.mem(), &hasher);
                mmr = mmr.apply_batch(&batch).unwrap();
            }
            let prune_loc = Location::new(50);
            let prune_pos = Position::try_from(prune_loc).unwrap();
            mmr = mmr.prune(prune_loc).await.unwrap();

            // Reopen below the previous in-memory boundary and verify the retained pinned nodes.
            let cap = mmr.leaves() - 80;
            _ = mmr.sync().await.unwrap();
            mmr = Mmr::init_at_most(context.child("cap"), &hasher, test_config(), cap)
                .await
                .unwrap();

            // A proof for the oldest item validates the pinned node boundary.
            mmr.proof(&hasher, prune_loc, 0).await.unwrap();
            // prune all remaining leaves 1 at a time.
            while mmr.size() > prune_pos {
                let cap = mmr.leaves() - 1;
                _ = mmr.sync().await.unwrap();
                mmr = Mmr::init_at_most(context.child("cap"), &hasher, test_config(), cap)
                    .await
                    .unwrap();
            }

            // Make sure pruning to an older location is a no-op.
            mmr = mmr.prune(prune_loc - 1).await.unwrap();
            assert_eq!(mmr.bounds().start, prune_loc);

            drop(mmr);
            assert!(matches!(
                Mmr::<_, Digest, Sequential>::init_at_most(
                    context.child("pruned_cap"),
                    &hasher,
                    test_config(),
                    prune_loc - 1
                )
                .await,
                Err(Error::ElementPruned(_))
            ));
        });
    }

    /// Create batch A, merkleize, create batch B via `merkleized_a.new_batch()`,
    /// merkleize, apply, and verify root matches a reference MMR.
    #[test_traced]
    fn test_full_mmr_batch_stacking() {
        let executor = deterministic::Runner::default();
        executor.start(|context| async move {
            let hasher: Standard<Sha256> = Standard::new(ForwardFold);

            // Build base full MMR with 10 elements.
            let mmr = Mmr::init(
                context.child("storage"),
                &Standard::<Sha256>::new(ForwardFold),
                test_config(),
            )
            .await
            .unwrap();

            let mut batch = mmr.new_batch();
            for i in 0u64..10 {
                let element = hasher.digest(&i.to_be_bytes());
                batch = batch.add(&hasher, &element);
            }
            let batch = batch.merkleize(mmr.mem(), &hasher);
            let mmr = mmr.apply_batch(&batch).unwrap();
            let mmr = mmr.sync().await.unwrap();

            // Batch A: add 5 elements.
            let mut batch_a = mmr.new_batch();
            for i in 10u64..15 {
                let element = hasher.digest(&i.to_be_bytes());
                batch_a = batch_a.add(&hasher, &element);
            }
            let merkleized_a = batch_a.merkleize(mmr.mem(), &hasher);

            // Batch B on merkleized A: add 5 more elements.
            let mut batch_b = merkleized_a.new_batch();
            for i in 15u64..20 {
                let element = hasher.digest(&i.to_be_bytes());
                batch_b = batch_b.add(&hasher, &element);
            }
            let merkleized_b = batch_b.merkleize(mmr.mem(), &hasher);
            let expected_root = merkleized_b.root(&hasher, 0).unwrap();

            // Apply.
            let mmr = mmr.apply_batch(&merkleized_b).unwrap();
            assert_eq!(mmr.root(&hasher, 0).unwrap(), expected_root);

            // Build a reference in-memory MMR with 20 elements to verify.
            let empty = mem::Mmr::new();
            let reference = build_test_mmr(&hasher, empty, 20);
            assert_eq!(
                mmr.root(&hasher, 0).unwrap(),
                reference.root(&hasher, 0).unwrap()
            );

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

    /// Regression: init_sync's "fresh start" path (retained tree ending at or before the sync
    /// start) resets the journal to the range start. A sync start past every stored blob leaves
    /// the journal unopened and checks the cleared journal. A sync start inside the stored blobs
    /// but past the recovered tree opens a journal too short to serve the range. Only that case
    /// reaches the raised journal_size. A stale size would initialize nodes_to_pin and the
    /// in-memory tree incorrectly or fail init.
    #[rstest]
    #[case::unopened(100, false)]
    #[case::opened_short(6, true)]
    #[test_traced]
    fn test_init_sync_fresh_start_updates_journal_size(
        #[case] leaves: usize,
        #[case] opened: bool,
    ) {
        let executor = deterministic::Runner::default();
        executor.start(|context| async move {
            let hasher = Standard::<Sha256>::new(ForwardFold);

            // Build an MMR with 5 leaves (size 8), sync, drop.
            let mmr =
                Mmr::<_, Digest, Sequential>::init(context.child("init"), &hasher, test_config())
                    .await
                    .unwrap();
            let mut batch = mmr.new_batch();
            for i in 0..5 {
                batch = batch.add(&hasher, &test_digest(i));
            }
            let batch = batch.merkleize(mmr.mem(), &hasher);
            let mmr = mmr.apply_batch(&batch).unwrap();
            let mmr = mmr.sync().await.unwrap();
            drop(mmr);

            // Build a reference MMR to the sync boundary to supply its pins and the expected root.
            let ref_cfg = Config {
                journal_partition: "ref-journal".into(),
                metadata_partition: "ref-metadata".into(),
                items_per_blob: NZU64!(7),
                write_buffer: NZUsize!(1024),
                replay_buffer: NZUsize!(1024),
                strategy: Sequential,
                page_cache: CacheRef::new(PAGE_SIZE, PAGE_CACHE_SIZE),
            };
            let ref_mmr =
                Mmr::<_, Digest, Sequential>::init(context.child("ref"), &hasher, ref_cfg)
                    .await
                    .unwrap();
            let mut batch = ref_mmr.new_batch();
            for i in 0..leaves {
                batch = batch.add(&hasher, &test_digest(i));
            }
            let batch = batch.merkleize(ref_mmr.mem(), &hasher);
            let ref_mmr = ref_mmr.apply_batch(&batch).unwrap();

            // Read the pins for a boundary at `leaves` from the reference, in the `nodes_to_pin`
            // order init_sync expects.
            let prune_loc = Location::from(leaves);
            let mut pinned = Vec::new();
            for pos in Family::nodes_to_pin(prune_loc) {
                pinned.push(ref_mmr.get_node(pos).await.unwrap().unwrap());
            }

            // Two local blobs with capacity 7 span positions 0..14 but recover a tree of size 8, so
            // both sync starts lie past the recovered tree and init_sync resets the journal to the
            // start. Position 197 (100 leaves) lies past every blob, so the journal stays unopened.
            // Position 10 (6 leaves) lies inside the second blob, so the journal is opened and
            // found too short. The asserts pin each case's start against 8 and 14.
            let expected_size = Position::try_from(prune_loc).unwrap();
            assert!(*expected_size > 8);
            assert_eq!(*expected_size < 14, opened);
            let sync_cfg = SyncConfig::<Digest, Sequential> {
                config: test_config(),
                range: non_empty_range!(prune_loc, Location::new(200)),
                pinned_nodes: Some(pinned),
            };
            let sync_mmr = Mmr::init_sync(context.child("sync"), sync_cfg)
                .await
                .unwrap();

            // The MMR should have size matching the prune boundary position.
            assert_eq!(sync_mmr.size(), expected_size);

            // Appending the same leaf to both trees must yield the same root. Below its start the
            // sync tree holds only the pins, so a match shows init_sync placed the pins and sized
            // the tree at the boundary.
            let batch = sync_mmr.new_batch().add(&hasher, &test_digest(999));
            let batch = batch.merkleize(sync_mmr.mem(), &hasher);
            let sync_mmr = sync_mmr.apply_batch(&batch).unwrap();
            let batch = ref_mmr.new_batch().add(&hasher, &test_digest(999));
            let batch = batch.merkleize(ref_mmr.mem(), &hasher);
            let ref_mmr = ref_mmr.apply_batch(&batch).unwrap();
            let root = ref_mmr.root(&hasher, 0).unwrap();
            assert_eq!(sync_mmr.root(&hasher, 0).unwrap(), root);
            ref_mmr.destroy().await.unwrap();

            // Ordinary init must recover the appended leaf from the reset journal and the
            // boundary pins from metadata.
            _ = sync_mmr.sync().await.unwrap();
            let mmr =
                Mmr::<_, Digest, Sequential>::init(context.child("reopen"), &hasher, test_config())
                    .await
                    .unwrap();
            assert_eq!(mmr.bounds(), prune_loc..prune_loc + 1);
            assert_eq!(mmr.root(&hasher, 0).unwrap(), root);

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