//! An _ordered_ variant of a [crate::qmdb::current] authenticated database for variable-size values
//!
//! This variant maintains the lexicographic-next active key for each active key, enabling exclusion
//! proofs (proving a key is currently inactive). Use [crate::qmdb::current::unordered::variable] if
//! exclusion proofs are not needed.
//!
//! See [Db] for the main database type and [super::ExclusionProof] for proving key inactivity.

pub use super::db::KeyValueProof;
use crate::{
    index::ordered::Index,
    journal::contiguous::variable::Journal,
    mmr::Location,
    qmdb::{
        any::{ordered::variable::Operation, value::VariableEncoding, VariableValue},
        current::VariableConfig as Config,
        operation::Key,
        Error,
    },
    translator::Translator,
};
use commonware_codec::{Codec, Read};
use commonware_cryptography::Hasher;
use commonware_runtime::{Clock, Metrics, Storage as RStorage};

pub type Db<E, K, V, H, T, const N: usize> =
    super::db::Db<E, Journal<E, Operation<K, V>>, K, VariableEncoding<V>, Index<T, Location>, H, N>;

impl<
        E: RStorage + Clock + Metrics,
        K: Key,
        V: VariableValue,
        H: Hasher,
        T: Translator,
        const N: usize,
    > Db<E, K, V, H, T, N>
where
    Operation<K, V>: Codec,
{
    /// Initializes a [Db] from the given `config`. Leverages parallel Merkleization to initialize
    /// the bitmap MMR if a thread pool is provided.
    pub async fn init(
        context: E,
        config: Config<T, <Operation<K, V> as Read>::Cfg>,
    ) -> Result<Self, Error> {
        crate::qmdb::current::init_variable(context, config, |ctx, t| Index::new(ctx, t)).await
    }
}

pub mod partitioned {
    //! A variant of [super] that uses a partitioned index for the snapshot.

    pub use super::KeyValueProof;
    use crate::{
        index::partitioned::ordered::Index,
        journal::contiguous::variable::Journal,
        mmr::Location,
        qmdb::{
            any::{
                ordered::variable::partitioned::Operation, value::VariableEncoding, VariableValue,
            },
            current::VariableConfig as Config,
            operation::Key,
            Error,
        },
        translator::Translator,
    };
    use commonware_codec::{Codec, Read};
    use commonware_cryptography::Hasher;
    use commonware_runtime::{Clock, Metrics, Storage as RStorage};

    /// A 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
    /// - `P = 3`: ~16 million partitions
    pub type Db<E, K, V, H, T, const P: usize, const N: usize> =
        crate::qmdb::current::ordered::db::Db<
            E,
            Journal<E, Operation<K, V>>,
            K,
            VariableEncoding<V>,
            Index<T, Location, P>,
            H,
            N,
        >;

    impl<
            E: RStorage + Clock + Metrics,
            K: Key,
            V: VariableValue,
            H: Hasher,
            T: Translator,
            const P: usize,
            const N: usize,
        > Db<E, K, V, H, T, P, N>
    where
        Operation<K, V>: Codec,
    {
        /// Initializes a [Db] from the given `config`. Leverages parallel Merkleization to initialize
        /// the bitmap MMR if a thread pool is provided.
        pub async fn init(
            context: E,
            config: Config<T, <Operation<K, V> as Read>::Cfg>,
        ) -> Result<Self, Error> {
            crate::qmdb::current::init_variable(context, config, |ctx, t| Index::new(ctx, t)).await
        }
    }
}

#[cfg(test)]
mod test {
    use crate::{
        kv::tests::{assert_gettable, assert_send},
        mmr::{hasher::Hasher as _, Location, StandardHasher},
        qmdb::{
            any::ordered::variable::Operation,
            current::{
                ordered::{db::KeyValueProof, variable::Db},
                proof::{OperationProof, RangeProof},
                tests::{apply_random_ops, variable_config},
            },
            store::{
                tests::{assert_log_store, assert_merkleized_store, assert_prunable_store},
                LogStore,
            },
            Error,
        },
        translator::OneCap,
    };
    use commonware_cryptography::{sha256::Digest, Digest as _, Hasher as _, Sha256};
    use commonware_macros::test_traced;
    use commonware_runtime::{deterministic, Runner as _};
    use commonware_utils::{bitmap::Prunable as BitMap, NZU64};
    use rand::RngCore;

    /// A type alias for the concrete [Db] type used in these unit tests.
    type CurrentTest = Db<deterministic::Context, Digest, Digest, Sha256, OneCap, 32>;

    /// Return a [Db] database initialized with a fixed config.
    async fn open_db(context: deterministic::Context, partition_prefix: String) -> CurrentTest {
        let cfg = variable_config::<OneCap>(&partition_prefix, &context);
        CurrentTest::init(context, cfg).await.unwrap()
    }

    /// Build a tiny database and make sure we can't convince the verifier that some old value of a
    /// key is active. We specifically test over the partial chunk case, since these bits are yet to
    /// be committed to the underlying MMR.
    #[test_traced("DEBUG")]
    pub fn test_current_db_verify_proof_over_bits_in_uncommitted_chunk() {
        let executor = deterministic::Runner::default();
        executor.start(|context| async move {
            let mut hasher = StandardHasher::<Sha256>::new();
            let partition = "build-small".into();
            let mut db = open_db(context, partition).await;

            // Add one key.
            let k = Sha256::fill(0x01);
            let v1 = Sha256::fill(0xA1);
            let finalized = {
                let mut batch = db.new_batch();
                batch.write(k, Some(v1));
                batch.merkleize(None).await.unwrap().finalize()
            };
            db.apply_batch(finalized).await.unwrap();

            let (_, op_loc) = db.any.get_with_loc(&k).await.unwrap().unwrap();
            let proof = db.key_value_proof(hasher.inner(), k).await.unwrap();

            // Proof should be verifiable against current root.
            let root = db.root();
            assert!(CurrentTest::verify_key_value_proof(
                hasher.inner(),
                k,
                v1,
                &proof,
                &root,
            ));

            let v2 = Sha256::fill(0xA2);
            // Proof should not verify against a different value.
            assert!(!CurrentTest::verify_key_value_proof(
                hasher.inner(),
                k,
                v2,
                &proof,
                &root,
            ));
            // Proof should not verify against a mangled next_key.
            let mut mangled_proof = proof.clone();
            mangled_proof.next_key = Sha256::fill(0xFF);
            assert!(!CurrentTest::verify_key_value_proof(
                hasher.inner(),
                k,
                v1,
                &mangled_proof,
                &root,
            ));

            // Update the key to a new value (v2), which inactivates the previous operation.
            let finalized = {
                let mut batch = db.new_batch();
                batch.write(k, Some(v2));
                batch.merkleize(None).await.unwrap().finalize()
            };
            db.apply_batch(finalized).await.unwrap();
            let root = db.root();

            // New value should not be verifiable against the old proof.
            assert!(!CurrentTest::verify_key_value_proof(
                hasher.inner(),
                k,
                v2,
                &proof,
                &root,
            ));

            // But the new value should verify against a new proof.
            let proof = db.key_value_proof(hasher.inner(), k).await.unwrap();
            assert!(CurrentTest::verify_key_value_proof(
                hasher.inner(),
                k,
                v2,
                &proof,
                &root,
            ));
            // Old value will not verify against new proof.
            assert!(!CurrentTest::verify_key_value_proof(
                hasher.inner(),
                k,
                v1,
                &proof,
                &root,
            ));

            // Create a proof of the now-inactive update operation assigning v1 to k against the
            // current root.
            let (p, _, chunks) = db
                .range_proof(hasher.inner(), op_loc, NZU64!(1))
                .await
                .unwrap();
            let proof_inactive = KeyValueProof {
                proof: OperationProof {
                    loc: op_loc,
                    chunk: chunks[0],
                    range_proof: p,
                },
                next_key: k,
            };
            // This proof should verify using verify_range_proof which does not check activity
            // status.
            let op = Operation::Update(crate::qmdb::any::ordered::Update {
                key: k,
                value: v1,
                next_key: k,
            });
            assert!(CurrentTest::verify_range_proof(
                hasher.inner(),
                &proof_inactive.proof.range_proof,
                proof_inactive.proof.loc,
                &[op],
                &[proof_inactive.proof.chunk],
                &root,
            ));
            // But this proof should *not* verify as a key value proof, since verification will see
            // that the operation is inactive.
            assert!(!CurrentTest::verify_key_value_proof(
                hasher.inner(),
                k,
                v1,
                &proof_inactive,
                &root,
            ));

            // Attempt #1 to "fool" the verifier:  change the location to that of an active
            // operation. This should not fool the verifier if we're properly validating the
            // inclusion of the operation itself, and not just the chunk.
            let (_, active_loc) = db.any.get_with_loc(&k).await.unwrap().unwrap();
            // The new location should differ but still be in the same chunk.
            assert_ne!(active_loc, proof_inactive.proof.loc);
            assert_eq!(
                BitMap::<32>::to_chunk_index(*active_loc),
                BitMap::<32>::to_chunk_index(*proof_inactive.proof.loc)
            );
            let mut fake_proof = proof_inactive.clone();
            fake_proof.proof.loc = active_loc;
            assert!(!CurrentTest::verify_key_value_proof(
                hasher.inner(),
                k,
                v1,
                &fake_proof,
                &root,
            ));

            // Attempt #2 to "fool" the verifier: Modify the chunk in the proof info to make it look
            // like the operation is active by flipping its corresponding bit to 1. This should not
            // fool the verifier if we are correctly incorporating the partial chunk information
            // into the root computation.
            let mut modified_chunk = proof_inactive.proof.chunk;
            let bit_pos = *proof_inactive.proof.loc;
            let byte_idx = bit_pos / 8;
            let bit_idx = bit_pos % 8;
            modified_chunk[byte_idx as usize] |= 1 << bit_idx;

            let mut fake_proof = proof_inactive.clone();
            fake_proof.proof.chunk = modified_chunk;
            assert!(!CurrentTest::verify_key_value_proof(
                hasher.inner(),
                k,
                v1,
                &fake_proof,
                &root,
            ));

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

    #[test_traced("DEBUG")]
    pub fn test_current_db_range_proofs() {
        let executor = deterministic::Runner::default();
        executor.start(|mut context| async move {
            let partition = "range-proofs".into();
            let mut hasher = StandardHasher::<Sha256>::new();
            let db = open_db(context.clone(), partition).await;
            let root = db.root();

            // Empty range proof should not crash or verify, since even an empty db has a single
            // commit op.
            let proof = RangeProof {
                proof: crate::mmr::Proof::default(),
                partial_chunk_digest: None,
                ops_root: Digest::EMPTY,
            };
            assert!(!CurrentTest::verify_range_proof(
                hasher.inner(),
                &proof,
                Location::new(0),
                &[],
                &[],
                &root,
            ));

            let rng_seed = context.next_u64();
            let mut db = apply_random_ops::<CurrentTest>(200, true, rng_seed, db)
                .await
                .unwrap();
            let finalized = db.new_batch().merkleize(None).await.unwrap().finalize();
            db.apply_batch(finalized).await.unwrap();
            let root = db.root();

            // Make sure size-constrained batches of operations are provable from the oldest
            // retained op to tip.
            let max_ops = 4;
            let end_loc = db.size().await;
            let start_loc = db.any.inactivity_floor_loc();

            for loc in *start_loc..*end_loc {
                let loc = Location::new(loc);
                let (proof, ops, chunks) = db
                    .range_proof(hasher.inner(), loc, NZU64!(max_ops))
                    .await
                    .unwrap();
                assert!(
                    CurrentTest::verify_range_proof(
                        hasher.inner(),
                        &proof,
                        loc,
                        &ops,
                        &chunks,
                        &root
                    ),
                    "failed to verify range at start_loc {start_loc}",
                );
                // Proof should not verify if we include extra chunks.
                let mut chunks_with_extra = chunks.clone();
                chunks_with_extra.push(chunks[chunks.len() - 1]);
                assert!(!CurrentTest::verify_range_proof(
                    hasher.inner(),
                    &proof,
                    loc,
                    &ops,
                    &chunks_with_extra,
                    &root,
                ));
            }

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

    #[test_traced("DEBUG")]
    pub fn test_current_db_key_value_proof() {
        let executor = deterministic::Runner::default();
        executor.start(|mut context| async move {
            let partition = "range-proofs".to_string();
            let mut hasher = StandardHasher::<Sha256>::new();
            let db = open_db(context.clone(), partition.clone()).await;
            let mut db = apply_random_ops::<CurrentTest>(500, true, context.next_u64(), db)
                .await
                .unwrap();
            let finalized = db.new_batch().merkleize(None).await.unwrap().finalize();
            db.apply_batch(finalized).await.unwrap();
            let root = db.root();

            // Confirm bad keys produce the expected error.
            let bad_key = Sha256::fill(0xAA);
            let res = db.key_value_proof(hasher.inner(), bad_key).await;
            assert!(matches!(res, Err(Error::KeyNotFound)));

            let start = *db.inactivity_floor_loc();
            for i in start..db.status.len() {
                if !db.status.get_bit(i) {
                    continue;
                }
                // Found an active operation! Create a proof for its active current key/value if
                // it's a key-updating operation.
                let op = db.any.log.read(Location::new(i)).await.unwrap();
                let (key, value) = match op {
                    Operation::Update(key_data) => (key_data.key, key_data.value),
                    Operation::CommitFloor(_, _) => continue,
                    _ => unreachable!("expected update or commit floor operation"),
                };
                let proof = db.key_value_proof(hasher.inner(), key).await.unwrap();

                // Proof should validate against the current value and correct root.
                assert!(CurrentTest::verify_key_value_proof(
                    hasher.inner(),
                    key,
                    value,
                    &proof,
                    &root
                ));
                // Proof should fail against the wrong value. Use hash instead of fill to ensure
                // the value differs from any key/value created by TestKey::from_seed (which uses
                // fill patterns).
                let wrong_val = Sha256::hash(&[0xFF]);
                assert!(!CurrentTest::verify_key_value_proof(
                    hasher.inner(),
                    key,
                    wrong_val,
                    &proof,
                    &root
                ));
                // Proof should fail against the wrong key.
                let wrong_key = Sha256::hash(&[0xEE]);
                assert!(!CurrentTest::verify_key_value_proof(
                    hasher.inner(),
                    wrong_key,
                    value,
                    &proof,
                    &root
                ));
                // Proof should fail against the wrong root.
                let wrong_root = Sha256::hash(&[0xDD]);
                assert!(!CurrentTest::verify_key_value_proof(
                    hasher.inner(),
                    key,
                    value,
                    &proof,
                    &wrong_root,
                ));
                // Proof should fail with the wrong next-key.
                let mut bad_proof = proof.clone();
                bad_proof.next_key = wrong_key;
                assert!(!CurrentTest::verify_key_value_proof(
                    hasher.inner(),
                    key,
                    value,
                    &bad_proof,
                    &root,
                ));
            }

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

    /// Repeatedly update the same key to a new value and ensure we can prove its current value
    /// after each update.
    #[test_traced("WARN")]
    pub fn test_current_db_proving_repeated_updates() {
        let executor = deterministic::Runner::default();
        executor.start(|context| async move {
            let mut hasher = StandardHasher::<Sha256>::new();
            let partition = "build-small".into();
            let mut db = open_db(context, partition).await;

            // Add one key.
            let k = Sha256::fill(0x00);
            let mut old_val = Sha256::fill(0x00);
            for i in 1u8..=255 {
                let v = Sha256::fill(i);
                let finalized = {
                    let mut batch = db.new_batch();
                    batch.write(k, Some(v));
                    batch.merkleize(None).await.unwrap().finalize()
                };
                db.apply_batch(finalized).await.unwrap();
                assert_eq!(db.get(&k).await.unwrap().unwrap(), v);
                let root = db.root();

                // Create a proof for the current value of k.
                let proof = db.key_value_proof(hasher.inner(), k).await.unwrap();
                assert!(
                    CurrentTest::verify_key_value_proof(hasher.inner(), k, v, &proof, &root),
                    "proof of update {i} failed to verify"
                );
                // Ensure the proof does NOT verify if we use the previous value.
                assert!(
                    !CurrentTest::verify_key_value_proof(hasher.inner(), k, old_val, &proof, &root),
                    "proof of update {i} verified when it should not have"
                );
                old_val = v;
            }

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

    /// Build a tiny database and confirm exclusion proofs work as expected with variable values.
    #[test_traced("DEBUG")]
    pub fn test_current_db_exclusion_proofs() {
        let executor = deterministic::Runner::default();
        executor.start(|context| async move {
            let mut hasher = StandardHasher::<Sha256>::new();
            let partition = "exclusion-proofs".into();
            let mut db = open_db(context, partition).await;

            let key_exists_1 = Sha256::fill(0x10);

            // We should be able to prove exclusion for any key against an empty db.
            let empty_root = db.root();
            let empty_proof = db
                .exclusion_proof(hasher.inner(), &key_exists_1)
                .await
                .unwrap();
            assert!(CurrentTest::verify_exclusion_proof(
                hasher.inner(),
                &key_exists_1,
                &empty_proof,
                &empty_root,
            ));

            // Add `key_exists_1` and test exclusion proving over the single-key database case.
            let v1 = Sha256::fill(0xA1);
            let finalized = {
                let mut batch = db.new_batch();
                batch.write(key_exists_1, Some(v1));
                batch.merkleize(None).await.unwrap().finalize()
            };
            db.apply_batch(finalized).await.unwrap();
            let root = db.root();

            // We shouldn't be able to generate an exclusion proof for a key already in the db.
            let result = db.exclusion_proof(hasher.inner(), &key_exists_1).await;
            assert!(matches!(result, Err(Error::KeyExists)));

            // Generate some valid exclusion proofs for keys on either side.
            let greater_key = Sha256::fill(0xFF);
            let lesser_key = Sha256::fill(0x00);
            let proof = db
                .exclusion_proof(hasher.inner(), &greater_key)
                .await
                .unwrap();
            let proof2 = db
                .exclusion_proof(hasher.inner(), &lesser_key)
                .await
                .unwrap();

            // Since there's only one span in the DB, the two exclusion proofs should be identical,
            // and the proof should verify any key but the one that exists in the db.
            assert_eq!(proof, proof2);
            // Any key except the one that exists should verify against this proof.
            assert!(CurrentTest::verify_exclusion_proof(
                hasher.inner(),
                &greater_key,
                &proof,
                &root,
            ));
            assert!(CurrentTest::verify_exclusion_proof(
                hasher.inner(),
                &lesser_key,
                &proof,
                &root,
            ));
            // Exclusion should fail if we test it on a key that exists.
            assert!(!CurrentTest::verify_exclusion_proof(
                hasher.inner(),
                &key_exists_1,
                &proof,
                &root,
            ));

            // Add a second key and test exclusion proving over the two-key database case.
            let key_exists_2 = Sha256::fill(0x30);
            let v2 = Sha256::fill(0xB2);

            let finalized = {
                let mut batch = db.new_batch();
                batch.write(key_exists_2, Some(v2));
                batch.merkleize(None).await.unwrap().finalize()
            };
            db.apply_batch(finalized).await.unwrap();
            let root = db.root();

            // Use a lesser/greater key that has a translated-key conflict based
            // on our use of OneCap translator.
            let lesser_key = Sha256::fill(0x0F); // < k1=0x10
            let greater_key = Sha256::fill(0x31); // > k2=0x30
            let middle_key = Sha256::fill(0x20); // between k1=0x10 and k2=0x30
            let proof = db
                .exclusion_proof(hasher.inner(), &greater_key)
                .await
                .unwrap();
            // Test the "cycle around" span. This should prove exclusion of greater_key & lesser
            // key, but fail on middle_key.
            assert!(CurrentTest::verify_exclusion_proof(
                hasher.inner(),
                &greater_key,
                &proof,
                &root,
            ));
            assert!(CurrentTest::verify_exclusion_proof(
                hasher.inner(),
                &lesser_key,
                &proof,
                &root,
            ));
            assert!(!CurrentTest::verify_exclusion_proof(
                hasher.inner(),
                &middle_key,
                &proof,
                &root,
            ));

            // Due to the cycle, lesser & greater keys should produce the same proof.
            let new_proof = db
                .exclusion_proof(hasher.inner(), &lesser_key)
                .await
                .unwrap();
            assert_eq!(proof, new_proof);

            // Test the inner span [k, k2).
            let proof = db
                .exclusion_proof(hasher.inner(), &middle_key)
                .await
                .unwrap();
            // `k` should fail since it's in the db.
            assert!(!CurrentTest::verify_exclusion_proof(
                hasher.inner(),
                &key_exists_1,
                &proof,
                &root,
            ));
            // `middle_key` should succeed since it's in range.
            assert!(CurrentTest::verify_exclusion_proof(
                hasher.inner(),
                &middle_key,
                &proof,
                &root,
            ));
            assert!(!CurrentTest::verify_exclusion_proof(
                hasher.inner(),
                &key_exists_2,
                &proof,
                &root,
            ));

            let conflicting_middle_key = Sha256::fill(0x11); // between k1=0x10 and k2=0x30
            assert!(CurrentTest::verify_exclusion_proof(
                hasher.inner(),
                &conflicting_middle_key,
                &proof,
                &root,
            ));

            // Using lesser/greater keys for the middle-proof should fail.
            assert!(!CurrentTest::verify_exclusion_proof(
                hasher.inner(),
                &greater_key,
                &proof,
                &root,
            ));
            assert!(!CurrentTest::verify_exclusion_proof(
                hasher.inner(),
                &lesser_key,
                &proof,
                &root,
            ));

            // Make the DB empty again by deleting the keys and check the empty case
            // again.
            let finalized = {
                let mut batch = db.new_batch();
                batch.write(key_exists_1, None);
                batch.write(key_exists_2, None);
                batch.merkleize(None).await.unwrap().finalize()
            };
            db.apply_batch(finalized).await.unwrap();
            db.sync().await.unwrap();
            let root = db.root();
            // This root should be different than the empty root from earlier since the DB now has a
            // non-zero number of operations.
            assert!(db.is_empty());
            assert_ne!(db.bounds().await.end, 0);
            assert_ne!(root, empty_root);

            let proof = db
                .exclusion_proof(hasher.inner(), &key_exists_1)
                .await
                .unwrap();
            assert!(CurrentTest::verify_exclusion_proof(
                hasher.inner(),
                &key_exists_1,
                &proof,
                &root,
            ));
            assert!(CurrentTest::verify_exclusion_proof(
                hasher.inner(),
                &key_exists_2,
                &proof,
                &root,
            ));

            // Try fooling the verifier with improper values.
            assert!(!CurrentTest::verify_exclusion_proof(
                hasher.inner(),
                &key_exists_1,
                &empty_proof, // wrong proof
                &root,
            ));
            assert!(!CurrentTest::verify_exclusion_proof(
                hasher.inner(),
                &key_exists_1,
                &proof,
                &empty_root, // wrong root
            ));
        });
    }

    #[allow(dead_code)]
    fn assert_db_futures_are_send(db: &mut CurrentTest, key: Digest, loc: Location) {
        assert_gettable(db, &key);
        assert_log_store(db);
        assert_prunable_store(db, loc);
        assert_merkleized_store(db, loc);
        assert_send(db.sync());
    }
}