// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed licenses.

//! Storage module for a auditable key directory

use crate::errors::StorageError;
use crate::storage::types::{DbRecord, StorageType};
use crate::{AkdLabel, AkdValue};

use async_trait::async_trait;
#[cfg(feature = "serde_serialization")]
use serde::{de::DeserializeOwned, Serialize};
use std::collections::HashMap;
use std::hash::Hash;
use std::marker::{Send, Sync};

pub mod cache;
pub mod transaction;
pub mod types;

/*
Various implementations supported by the library are imported here and usable at various checkpoints
*/
pub mod manager;
pub mod memory;

pub use manager::StorageManager;

#[cfg(any(test, feature = "public_tests"))]
pub mod tests;

/// Denotes the "state" when a batch_set is being called in the data layer
pub enum DbSetState {
    /// Being called as part of a transaction commit operation
    TransactionCommit,
    /// Being called as a general, in-line operation
    General,
}

/// Storable represents an _item_ which can be stored in the storage layer
#[cfg(feature = "serde_serialization")]
pub trait Storable: Clone + Serialize + DeserializeOwned + Sync + 'static {
    /// This particular storage will have a key type
    type StorageKey: Clone + Serialize + Eq + Hash + Send + Sync + std::fmt::Debug;

    /// Must return a valid storage type
    fn data_type() -> StorageType;

    /// Retrieve an instance of the id of this storable. The combination of the
    /// storable's StorageType and this id are _globally_ unique
    fn get_id(&self) -> Self::StorageKey;

    /// Retrieve the full binary version of a key (for comparisons)
    fn get_full_binary_id(&self) -> Vec<u8> {
        Self::get_full_binary_key_id(&self.get_id())
    }

    /// Retrieve the full binary version of a key (for comparisons)
    fn get_full_binary_key_id(key: &Self::StorageKey) -> Vec<u8>;

    /// Reformat a key from the full-binary specification
    fn key_from_full_binary(bin: &[u8]) -> Result<Self::StorageKey, String>;
}

/// Storable represents an _item_ which can be stored in the storage layer
#[cfg(not(feature = "serde_serialization"))]
pub trait Storable: Clone + Sync + 'static {
    /// This particular storage will have a key type
    type StorageKey: Clone + Eq + Hash + Send + Sync + std::fmt::Debug;

    /// Must return a valid storage type
    fn data_type() -> StorageType;

    /// Retrieve an instance of the id of this storable. The combination of the
    /// storable's StorageType and this id are _globally_ unique
    fn get_id(&self) -> Self::StorageKey;

    /// Retrieve the full binary version of a key (for comparisons)
    fn get_full_binary_id(&self) -> Vec<u8> {
        Self::get_full_binary_key_id(&self.get_id())
    }

    /// Retrieve the full binary version of a key (for comparisons)
    fn get_full_binary_key_id(key: &Self::StorageKey) -> Vec<u8>;

    /// Reformat a key from the full-binary specification
    fn key_from_full_binary(bin: &[u8]) -> Result<Self::StorageKey, String>;
}

/// A database implementation backing storage for the AKD
#[async_trait]
pub trait Database: Send + Sync {
    /// Set a record in the database
    async fn set(&self, record: DbRecord) -> Result<(), StorageError>;

    /// Set multiple records in the database with a minimal set of operations
    async fn batch_set(
        &self,
        records: Vec<DbRecord>,
        state: DbSetState,
    ) -> Result<(), StorageError>;

    /// Retrieve a stored record from the database
    async fn get<St: Storable>(&self, id: &St::StorageKey) -> Result<DbRecord, StorageError>;

    /// Retrieve a batch of records by id from the database
    async fn batch_get<St: Storable>(
        &self,
        ids: &[St::StorageKey],
    ) -> Result<Vec<DbRecord>, StorageError>;

    /* User data searching */

    /// Retrieve the user data for a given user
    async fn get_user_data(&self, username: &AkdLabel) -> Result<types::KeyData, StorageError>;

    /// Retrieve a specific state for a given user
    async fn get_user_state(
        &self,
        username: &AkdLabel,
        flag: types::ValueStateRetrievalFlag,
    ) -> Result<types::ValueState, StorageError>;

    /// Retrieve the user -> state version mapping in bulk. This is the same as get_user_states but with less data retrieved from the storage layer
    async fn get_user_state_versions(
        &self,
        usernames: &[AkdLabel],
        flag: types::ValueStateRetrievalFlag,
    ) -> Result<HashMap<AkdLabel, (u64, AkdValue)>, StorageError>;
}

/// Optional storage layer utility functions for debug and test purposes
#[async_trait]
pub trait StorageUtil: Database {
    /// Retrieves all stored records of a given type from the data layer, ignoring any caching or transaction pending
    async fn batch_get_type_direct<St: Storable>(&self) -> Result<Vec<DbRecord>, StorageError>;

    /// Retrieves all stored records from the data layer, ignoring any caching or transaction pending
    async fn batch_get_all_direct(&self) -> Result<Vec<DbRecord>, StorageError>;
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed licenses.

//! Test utilities of storage layers implementing the storage primatives for AKD

use crate::errors::StorageError;
use crate::storage::types::*;
use crate::storage::Database;
use crate::storage::StorageManager;
use crate::tree_node::*;
use crate::utils::byte_arr_from_u64;
use crate::NodeLabel;
use crate::{AkdLabel, AkdValue};

use akd_core::hash::EMPTY_DIGEST;
use akd_core::AzksValue;
use rand::distributions::Alphanumeric;
use rand::{thread_rng, Rng};
use std::time::{Duration, Instant};

type Azks = crate::append_only_zks::Azks;
type TreeNode = crate::tree_node::TreeNode;
type PvTreeNode = crate::tree_node::TreeNodeWithPreviousValue;

// *** Tests *** //

#[cfg(test)]
mod memory_storage_tests {
    use crate::storage::memory::AsyncInMemoryDatabase;
    use serial_test::serial;

    #[tokio::test]
    #[serial]
    async fn test_in_memory_db() {
        let db = AsyncInMemoryDatabase::new();
        crate::storage::tests::run_test_cases_for_storage_impl(db).await;
    }
}

// *** Run the test cases for a given data-layer impl *** //
/// Run the storage-layer test suite for a given storage implementation.
/// This is public because it can be used by other implemented storage layers
/// for consistency checks (e.g. mysql, memcached, etc)
pub async fn run_test_cases_for_storage_impl<S: Database>(db: S) -> StorageManager<S> {
    test_get_and_set_item(&db).await;
    test_user_data(&db).await;
    test_batch_get_items(&db).await;

    let manager = StorageManager::new_no_cache(db);
    test_transactions(&manager).await;
    test_tombstoning_data(&manager).await.unwrap();
    manager
}

// *** New Test Helper Functions *** //
async fn test_get_and_set_item<Ns: Database>(storage: &Ns) {
    // === Azks storage === //
    let azks = Azks {
        latest_epoch: 34,
        num_nodes: 10,
    };

    let set_result = storage.set(DbRecord::Azks(azks.clone())).await;
    assert_eq!(Ok(()), set_result);

    let get_result = storage
        .get::<Azks>(&crate::append_only_zks::DEFAULT_AZKS_KEY)
        .await;
    if let Ok(DbRecord::Azks(got_azks)) = get_result {
        assert_eq!(got_azks.latest_epoch, azks.latest_epoch);
        assert_eq!(got_azks.num_nodes, azks.num_nodes);
    } else {
        panic!("Failed to retrieve AZKS");
    }

    // === TreeNode storage === //

    let node = TreeNode {
        label: NodeLabel::new(byte_arr_from_u64(13), 4),
        last_epoch: 34,
        min_descendant_epoch: 1,
        parent: NodeLabel::new(byte_arr_from_u64(1), 1),
        node_type: TreeNodeType::Leaf,
        left_child: None,
        right_child: None,
        hash: AzksValue(EMPTY_DIGEST),
    };
    let mut node2 = node.clone();
    node2.label = NodeLabel::new(byte_arr_from_u64(16), 4);

    let key = NodeKey(NodeLabel::new(byte_arr_from_u64(13), 4));
    let key2 = NodeKey(NodeLabel::new(byte_arr_from_u64(16), 4));

    let set_result = storage
        .set(DbRecord::TreeNode(PvTreeNode::from_tree_node(node.clone())))
        .await;
    assert_eq!(Ok(()), set_result);

    let set_result = storage
        .set(DbRecord::TreeNode(PvTreeNode::from_tree_node(
            node2.clone(),
        )))
        .await;
    assert_eq!(Ok(()), set_result);

    let get_result = storage.get::<PvTreeNode>(&key).await;
    if let Ok(DbRecord::TreeNode(got_node)) = get_result {
        assert_eq!(got_node.label, node.label);
        assert_eq!(got_node.latest_node.parent, node.parent);
        assert_eq!(got_node.latest_node.node_type, node.node_type);
        assert_eq!(got_node.latest_node.last_epoch, node.last_epoch);
    } else {
        panic!("Failed to retrieve History Tree Node");
    }

    let get_result = storage.get::<PvTreeNode>(&key2).await;
    if let Err(err) = get_result {
        panic!("Failed to retrieve history tree node (2) {:?}", err)
    }

    // === ValueState storage === //
    let key = ValueStateKey("test".as_bytes().to_vec(), 1);
    let value = ValueState {
        username: AkdLabel::from("test"),
        epoch: 1,
        label: NodeLabel::new(byte_arr_from_u64(1), 1),
        version: 1,
        value: AkdValue::from("abc123"),
    };
    let set_result = storage.set(DbRecord::ValueState(value.clone())).await;
    assert_eq!(Ok(()), set_result);

    let get_result = storage.get::<ValueState>(&key).await;
    if let Ok(DbRecord::ValueState(got_state)) = get_result {
        assert_eq!(got_state.username, value.username);
        assert_eq!(got_state.epoch, value.epoch);
        assert_eq!(got_state.label, value.label);
        assert_eq!(got_state.value, value.value);
        assert_eq!(got_state.version, value.version);
    } else {
        panic!("Failed to retrieve history node state");
    }
}

async fn test_batch_get_items<Ns: Database>(storage: &Ns) {
    let mut rand_users: Vec<Vec<u8>> = vec![];
    for _ in 0..20 {
        let str: String = thread_rng()
            .sample_iter(&Alphanumeric)
            .take(30)
            .map(char::from)
            .collect();
        rand_users.push(str.as_bytes().to_vec());
    }

    let mut data = Vec::new();

    let mut epoch = 1;
    for value in rand_users.iter() {
        for user in rand_users.iter() {
            data.push(DbRecord::ValueState(ValueState {
                value: AkdValue(value.clone()),
                version: epoch,
                label: NodeLabel {
                    label_val: byte_arr_from_u64(1),
                    label_len: 1u32,
                },
                epoch,
                username: AkdLabel(user.clone()),
            }));
        }
        epoch += 1;
    }

    let tic = Instant::now();
    assert_eq!(
        Ok(()),
        storage
            .batch_set(data.clone(), crate::storage::DbSetState::General)
            .await
    );
    let toc: Duration = Instant::now() - tic;
    println!("Storage batch op: {} ms", toc.as_millis());
    let got = storage
        .get::<ValueState>(&ValueStateKey(rand_users[0].clone(), 10))
        .await;
    if got.is_err() {
        panic!("Failed to retrieve a user after batch insert");
    }

    let keys: Vec<ValueStateKey> = rand_users
        .iter()
        .map(|user| ValueStateKey(user.clone(), 1))
        .collect();
    let got_all = storage.batch_get::<ValueState>(&keys).await;
    match got_all {
        Err(_) => panic!("Failed to retrieve batch of user at specific epochs"),
        Ok(lst) if lst.len() != rand_users.len() => {
            panic!(
                "Retrieved list length does not match input length {} != {}",
                lst.len(),
                rand_users.len()
            );
        }
        Ok(results) => {
            // correct length, now check the values
            for result in results.into_iter() {
                // find the initial record with the same username & epoch
                let initial_record = data
                    .iter()
                    .find(|&x| {
                        if let DbRecord::ValueState(value_state) = &x {
                            if let DbRecord::ValueState(retrieved_state) = &result {
                                return value_state.username == retrieved_state.username
                                    && value_state.epoch == retrieved_state.epoch;
                            }
                        }
                        false
                    })
                    .cloned();
                // assert it matches what was given matches what was retrieved
                assert_eq!(Some(result), initial_record);
            }
        }
    }

    let user_keys: Vec<_> = rand_users
        .iter()
        .map(|user| AkdLabel(user.clone()))
        .collect();
    let got_all_min_states = storage
        .get_user_state_versions(&user_keys, ValueStateRetrievalFlag::MinEpoch)
        .await;
    // should be the same thing as the previous get
    match got_all_min_states {
        Err(err) => panic!("Failed to retrieve batch of user at min epochs: {:?}", err),
        Ok(lst) if lst.len() != rand_users.len() => {
            panic!(
                "Retrieved list length does not match input length {} != {}",
                lst.len(),
                rand_users.len()
            );
        }
        Ok(results) => {
            // correct length, now check the values
            for result in results.into_iter() {
                // find the initial record with the same username & epoch
                let initial_record = data
                    .iter()
                    .find(|&x| {
                        if let DbRecord::ValueState(value_state) = &x {
                            return value_state.username == result.0
                                && value_state.version == result.1 .0;
                        }
                        false
                    })
                    .cloned()
                    .map(|item| {
                        if let DbRecord::ValueState(value_state) = &item {
                            value_state.version
                        } else {
                            0u64
                        }
                    });

                // assert it matches what was given matches what was retrieved
                assert_eq!(Some(result.1 .0), initial_record);
            }
        }
    }

    let got_all_max_states = storage
        .get_user_state_versions(&user_keys, ValueStateRetrievalFlag::MaxEpoch)
        .await;
    // should be the same thing as the previous get
    match got_all_max_states {
        Err(err) => panic!("Failed to retrieve batch of user at min epochs: {:?}", err),
        Ok(lst) if lst.len() != rand_users.len() => {
            panic!(
                "Retrieved list length does not match input length {} != {}",
                lst.len(),
                rand_users.len()
            );
        }
        Ok(results) => {
            // correct length, now check the values
            for result in results.into_iter() {
                // find the initial record with the same username & epoch
                let initial_record = data
                    .iter()
                    .find(|&x| {
                        if let DbRecord::ValueState(value_state) = &x {
                            return value_state.username == result.0
                                && value_state.version == result.1 .0;
                        }
                        false
                    })
                    .cloned()
                    .map(|item| {
                        if let DbRecord::ValueState(value_state) = &item {
                            value_state.version
                        } else {
                            0u64
                        }
                    });
                // assert it matches what was given matches what was retrieved
                assert_eq!(Some(result.1 .0), initial_record);
            }
        }
    }
}

async fn test_transactions<S: Database>(storage: &StorageManager<S>) {
    let mut rand_users: Vec<Vec<u8>> = vec![];
    for _ in 0..20 {
        let str: String = thread_rng()
            .sample_iter(&Alphanumeric)
            .take(30)
            .map(char::from)
            .collect();
        rand_users.push(str.as_bytes().to_vec());
    }

    let mut data = Vec::new();

    let mut epoch = 1;
    for value in rand_users.iter() {
        for user in rand_users.iter() {
            data.push(DbRecord::ValueState(ValueState {
                value: AkdValue(value.clone()),
                version: 1u64,
                label: NodeLabel {
                    label_val: byte_arr_from_u64(1),
                    label_len: 1u32,
                },
                epoch,
                username: AkdLabel(user.clone()),
            }));
        }
        epoch += 1;
    }

    data.push(DbRecord::Azks(Azks {
        latest_epoch: 1,
        num_nodes: 34,
    }));

    let new_data = data
        .iter()
        .map(|item| {
            let new_item = item.clone();
            match &item {
                DbRecord::ValueState(new_state) => {
                    let mut copied_state = new_state.clone();
                    copied_state.epoch += 10000;
                    DbRecord::ValueState(copied_state)
                }
                DbRecord::Azks(azks) => DbRecord::Azks(Azks {
                    latest_epoch: azks.latest_epoch + 10000,
                    num_nodes: azks.num_nodes,
                }),
                _ => new_item,
            }
        })
        .collect();

    let tic = Instant::now();
    assert_eq!(Ok(()), storage.batch_set(data).await);
    let toc: Duration = Instant::now() - tic;
    println!("Storage batch op: {} ms", toc.as_millis());
    let got = storage
        .get::<ValueState>(&ValueStateKey(rand_users[0].clone(), 10))
        .await;
    if got.is_err() {
        panic!("Failed to retrieve a user after batch insert");
    }

    let tic = Instant::now();
    assert!(storage.begin_transaction());
    assert_eq!(Ok(()), storage.batch_set(new_data).await);
    assert!(storage.commit_transaction().await.is_ok());
    let toc: Duration = Instant::now() - tic;
    println!("Transactional storage batch op: {} ms", toc.as_millis());

    let got = storage
        .get::<ValueState>(&ValueStateKey(rand_users[0].clone(), 10 + 10000))
        .await;
    if got.is_err() {
        panic!("Failed to retrieve a user after batch insert");
    }
}

async fn test_user_data<S: Database>(storage: &S) {
    let rand_user = thread_rng()
        .sample_iter(&Alphanumeric)
        .take(30)
        .map(char::from)
        .collect::<String>()
        .as_bytes()
        .to_vec();
    let rand_value = thread_rng()
        .sample_iter(&Alphanumeric)
        .take(1028)
        .map(char::from)
        .collect::<String>()
        .as_bytes()
        .to_vec();
    let mut sample_state = ValueState {
        value: AkdValue(rand_value.clone()),
        version: 1u64,
        label: NodeLabel {
            label_val: byte_arr_from_u64(1),
            label_len: 1u32,
        },
        epoch: 1u64,
        username: AkdLabel(rand_user),
    };
    let mut sample_state_2 = sample_state.clone();
    sample_state_2.username = AkdLabel::from("test_user");

    let result = storage
        .set(DbRecord::ValueState(sample_state.clone()))
        .await;
    assert_eq!(Ok(()), result);

    sample_state.version = 2u64;
    sample_state.epoch = 123u64;
    let result = storage
        .set(DbRecord::ValueState(sample_state.clone()))
        .await;
    assert_eq!(Ok(()), result);

    sample_state.version = 3u64;
    sample_state.epoch = 456u64;
    let result = storage
        .set(DbRecord::ValueState(sample_state.clone()))
        .await;
    assert_eq!(Ok(()), result);

    let data = storage.get_user_data(&sample_state.username).await.unwrap();
    assert_eq!(3, data.states.len());

    let versions = data
        .states
        .into_iter()
        .map(|state| state.version)
        .collect::<Vec<_>>();
    assert_eq!(vec![1, 2, 3], versions);

    // At this point the DB has structure (for MySQL):
    /*
    mysql> USE default;
    Reading table information for completion of table and column names
    You can turn off this feature to get a quicker startup with -A

    Database changed
    mysql> SHOW TABLES;
    +-------------------+
    | Tables_in_default |
    +-------------------+
    | data              |
    | user_data         |
    +-------------------+
    2 rows in set (0.00 sec)

    mysql> SELECT * FROM user_data;
    +--------------------------------+-------+---------+----------------+----------------+-------------------+
    | username                       | epoch | version | node_label_val | node_label_len | data              |
    +--------------------------------+-------+---------+----------------+----------------+-------------------+
    | do3zfiXa0IUKznscp06jtc6KfHJudy |     1 |       1 |              1 |              1 | 8owmLSoZi...B9pu8 |
    | do3zfiXa0IUKznscp06jtc6KfHJudy |   123 |       2 |              1 |              1 | 8owmLSoZi...B9pu8 |
    | do3zfiXa0IUKznscp06jtc6KfHJudy |   456 |       3 |              1 |              1 | 8owmLSoZi...B9pu8 |
    +--------------------------------+-------+---------+----------------+----------------+-------------------+
    3 rows in set (0.00 sec)
    */

    let specific_result = storage
        .get_user_state(
            &sample_state.username,
            ValueStateRetrievalFlag::SpecificVersion(2),
        )
        .await;
    assert_eq!(
        Ok(ValueState {
            epoch: 123,
            version: 2,
            label: NodeLabel::new(byte_arr_from_u64(1), 1),
            value: AkdValue(rand_value.clone()),
            username: sample_state.username.clone(),
        }),
        specific_result
    );

    let specifc_result = storage
        .get::<ValueState>(&ValueStateKey(sample_state.username.to_vec(), 123))
        .await;
    if let Ok(DbRecord::ValueState(state)) = specifc_result {
        assert_eq!(
            ValueState {
                epoch: 123,
                version: 2,
                label: NodeLabel::new(byte_arr_from_u64(1), 1),
                value: AkdValue(rand_value.clone()),
                username: sample_state.username.clone(),
            },
            state
        );
    } else {
        panic!("Unable to retrieve user state object");
    }

    let missing_result = storage
        .get_user_state(
            &sample_state.username,
            ValueStateRetrievalFlag::SpecificVersion(100),
        )
        .await;
    assert!(matches!(missing_result, Err(StorageError::NotFound(_)),));

    let specific_result = storage
        .get_user_state(
            &sample_state.username,
            ValueStateRetrievalFlag::SpecificEpoch(123),
        )
        .await;
    assert_eq!(
        Ok(ValueState {
            epoch: 123,
            version: 2,
            label: NodeLabel::new(byte_arr_from_u64(1), 1),
            value: AkdValue(rand_value.clone()),
            username: sample_state.username.clone(),
        }),
        specific_result
    );

    let specific_result = storage
        .get_user_state(&sample_state.username, ValueStateRetrievalFlag::MinEpoch)
        .await;
    assert_eq!(
        Ok(ValueState {
            epoch: 1,
            version: 1,
            label: NodeLabel::new(byte_arr_from_u64(1), 1),
            value: AkdValue(rand_value.clone()),
            username: sample_state.username.clone(),
        }),
        specific_result
    );

    let specific_result = storage
        .get_user_state(&sample_state.username, ValueStateRetrievalFlag::MaxEpoch)
        .await;
    assert_eq!(
        Ok(ValueState {
            epoch: 456,
            version: 3,
            label: NodeLabel::new(byte_arr_from_u64(1), 1),
            value: AkdValue(rand_value.clone()),
            username: sample_state.username.clone(),
        }),
        specific_result
    );

    // Vector operations

    let mut vector_of_states = vec![sample_state_2.clone()];
    sample_state_2.version = 2;
    sample_state_2.epoch = 234;
    vector_of_states.push(sample_state_2.clone());

    sample_state_2.version = 3;
    sample_state_2.epoch = 345;
    vector_of_states.push(sample_state_2.clone());
    sample_state_2.version = 4;
    sample_state_2.epoch = 456;
    vector_of_states.push(sample_state_2.clone());

    let records = vector_of_states
        .into_iter()
        .map(DbRecord::ValueState)
        .collect::<Vec<_>>();
    let result = storage
        .batch_set(records, crate::storage::DbSetState::General)
        .await;
    assert_eq!(Ok(()), result);

    let data = storage.get_user_data(&sample_state_2.username).await;
    assert_eq!(4, data.unwrap().states.len());
}

async fn test_tombstoning_data<S: Database>(
    storage: &StorageManager<S>,
) -> Result<(), crate::errors::AkdError> {
    let rand_user = thread_rng()
        .sample_iter(&Alphanumeric)
        .take(30)
        .map(char::from)
        .collect::<String>()
        .as_bytes()
        .to_vec();
    let rand_value = rand_user.clone();

    let mut sample_state = ValueState {
        value: AkdValue(rand_value.clone()),
        version: 1u64,
        label: NodeLabel {
            label_val: byte_arr_from_u64(1),
            label_len: 1u32,
        },
        epoch: 1u64,
        username: AkdLabel(rand_user.clone()),
    };
    let mut sample_state2 = sample_state.clone();
    sample_state2.username = AkdLabel::from("tombstone_test_user");

    // Load up a bunch of data into the storage layer
    for i in 0..5 {
        sample_state.version = i;
        sample_state.epoch = i;
        sample_state2.version = i;
        sample_state2.epoch = i;

        assert_eq!(
            Ok(()),
            storage
                .set(DbRecord::ValueState(sample_state.clone()))
                .await
        );
        assert_eq!(
            Ok(()),
            storage
                .set(DbRecord::ValueState(sample_state2.clone()))
                .await
        );
    }

    let data = storage.get_user_data(&sample_state.username).await.unwrap();
    assert_eq!(5, data.states.len());
    let data = storage
        .get_user_data(&sample_state2.username)
        .await
        .unwrap();
    assert_eq!(5, data.states.len());

    let keys_to_tombstone = [
        ValueStateKey("tombstone_test_user".as_bytes().to_vec(), 0u64),
        ValueStateKey("tombstone_test_user".as_bytes().to_vec(), 1u64),
        ValueStateKey(sample_state.username.to_vec(), 0u64),
        ValueStateKey(sample_state.username.to_vec(), 1u64),
        ValueStateKey(sample_state.username.to_vec(), 2u64),
    ];

    // tombstone the given states
    storage.tombstone_value_states(&keys_to_tombstone).await?;

    for label in [AkdLabel::from("tombstone_test_user"), AkdLabel(rand_user)] {
        for version in 0..5 {
            let key = ValueStateKey(label.to_vec(), version);
            let got = storage.get::<ValueState>(&key).await?;

            if let DbRecord::ValueState(value_state) = got {
                assert_eq!(version, value_state.epoch);
                if keys_to_tombstone.contains(&key) {
                    // should be a tombstone
                    assert_eq!(crate::TOMBSTONE.to_vec(), value_state.value.0);
                } else {
                    // should NOT be a tombstone
                    assert_ne!(crate::TOMBSTONE.to_vec(), value_state.value.0);
                }
            } else {
                panic!("Unable to retrieve value state {:?}", key);
            }
        }
    }
    Ok(())
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed licenses.

//! A simple in-memory transaction object to minimize data-layer operations

use crate::errors::StorageError;
use crate::storage::types::DbRecord;
use crate::storage::types::ValueState;
use crate::storage::types::ValueStateRetrievalFlag;
use crate::storage::Storable;

use dashmap::DashMap;
#[cfg(feature = "runtime_metrics")]
use log::{debug, error, info, trace, warn};
use std::collections::HashMap;
#[cfg(feature = "runtime_metrics")]
use std::sync::atomic::AtomicU64;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;

/// Represents an in-memory transaction, keeping a mutable state
/// of the changes. When you "commit" this transaction, you return the
/// collection of values which need to be written to the storage layer
/// including all mutations. Rollback simply empties the transaction state.
#[derive(Clone)]
pub struct Transaction {
    mods: Arc<DashMap<Vec<u8>, DbRecord>>,
    active: Arc<AtomicBool>,

    #[cfg(feature = "runtime_metrics")]
    num_reads: Arc<AtomicU64>,
    #[cfg(feature = "runtime_metrics")]
    num_writes: Arc<AtomicU64>,
}

unsafe impl Send for Transaction {}
unsafe impl Sync for Transaction {}

impl std::fmt::Debug for Transaction {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "a lone transaction")
    }
}

impl Default for Transaction {
    fn default() -> Self {
        Self::new()
    }
}

impl Transaction {
    /// Instantiate a new transaction instance
    pub fn new() -> Self {
        Self {
            mods: Arc::new(DashMap::new()),
            active: Arc::new(AtomicBool::new(false)),

            #[cfg(feature = "runtime_metrics")]
            num_reads: Arc::new(AtomicU64::new(0)),
            #[cfg(feature = "runtime_metrics")]
            num_writes: Arc::new(AtomicU64::new(0)),
        }
    }

    /// Get the current number of items currently in the transaction modifications set
    pub fn count(&self) -> usize {
        self.mods.len()
    }

    /// Log metrics about the current transaction instance. Metrics will be cleared after log call
    pub fn log_metrics(&self, _level: log::Level) {
        #[cfg(feature = "runtime_metrics")]
        {
            let r = self.num_reads.swap(0, Ordering::Relaxed);
            let w = self.num_writes.swap(0, Ordering::Relaxed);

            let msg = format!("Transaction writes: {w}, Transaction reads: {r}");

            match _level {
                log::Level::Trace => trace!("{}", msg),
                log::Level::Debug => debug!("{}", msg),
                log::Level::Info => info!("{}", msg),
                log::Level::Warn => warn!("{}", msg),
                _ => error!("{}", msg),
            }
        }
    }

    /// Start a transaction in the storage layer
    pub fn begin_transaction(&self) -> bool {
        !self.active.swap(true, Ordering::Relaxed)
    }

    /// Commit a transaction in the storage layer
    pub fn commit_transaction(&self) -> Result<Vec<DbRecord>, StorageError> {
        if !self.active.load(Ordering::Relaxed) {
            return Err(StorageError::Transaction(
                "Transaction not currently active".to_string(),
            ));
        }

        // copy all the updated values out
        let mut records = self
            .mods
            .iter()
            .map(|p| p.value().clone())
            .collect::<Vec<_>>();

        // sort according to transaction priority
        records.sort_by_key(|r| r.transaction_priority());

        // flush the trans log
        self.mods.clear();

        self.active.store(false, Ordering::Relaxed);
        Ok(records)
    }

    /// Rollback a transaction
    pub fn rollback_transaction(&self) -> Result<(), StorageError> {
        if !self.active.load(Ordering::Relaxed) {
            return Err(StorageError::Transaction(
                "Transaction not currently active".to_string(),
            ));
        }

        // rollback
        self.mods.clear();

        self.active.store(false, Ordering::Relaxed);
        Ok(())
    }

    /// Retrieve a flag determining if there is a transaction active
    pub fn is_transaction_active(&self) -> bool {
        self.active.load(Ordering::Relaxed)
    }

    /// Hit test the current transaction to see if it is currently active
    pub fn get<St: Storable>(&self, key: &St::StorageKey) -> Option<DbRecord> {
        let bin_id = St::get_full_binary_key_id(key);

        let out = self.mods.get(&bin_id).map(|p| p.value().clone());
        #[cfg(feature = "runtime_metrics")]
        if out.is_some() {
            self.num_reads.fetch_add(1, Ordering::Relaxed);
        }
        out
    }

    /// Set a batch of values into the cache
    pub fn batch_set(&self, records: &[DbRecord]) {
        for record in records {
            self.mods
                .insert(record.get_full_binary_id(), record.clone());
        }

        #[cfg(feature = "runtime_metrics")]
        {
            self.num_writes.fetch_add(1, Ordering::Relaxed);
        }
    }

    /// Set a value in the transaction to be committed at transaction commit time
    pub fn set(&self, record: &DbRecord) {
        let bin_id = record.get_full_binary_id();

        self.mods.insert(bin_id, record.clone());

        #[cfg(feature = "runtime_metrics")]
        {
            self.num_writes.fetch_add(1, Ordering::Relaxed);
        }
    }

    /// Retrieve all of the user data for a given username
    ///
    /// Note: This is a FULL SCAN operation of the entire transaction log
    pub fn get_users_data(
        &self,
        usernames: &[crate::AkdLabel],
    ) -> HashMap<crate::AkdLabel, Vec<ValueState>> {
        let mut results: HashMap<crate::AkdLabel, Vec<ValueState>> = HashMap::new();

        let mut set = std::collections::HashSet::with_capacity(usernames.len());
        for username in usernames.iter() {
            if !set.contains(username) {
                set.insert(username.clone());
            }
        }

        for pair in self.mods.iter() {
            if let DbRecord::ValueState(value_state) = pair.value() {
                if set.contains(&value_state.username) {
                    if results.contains_key(&value_state.username) {
                        if let Some(item) = results.get_mut(&value_state.username) {
                            item.push(value_state.clone())
                        }
                    } else {
                        results.insert(value_state.username.clone(), vec![value_state.clone()]);
                    }
                }
            }
        }

        // sort all the value lists by epoch
        for (_k, v) in results.iter_mut() {
            v.sort_unstable_by(|a, b| a.epoch.cmp(&b.epoch));
        }

        results
    }

    /// Retrieve the user state given the specified value state retrieval mode
    ///
    /// Note: This is a FULL SCAN operation of the entire transaction log
    pub fn get_user_state(
        &self,
        username: &crate::AkdLabel,
        flag: ValueStateRetrievalFlag,
    ) -> Option<ValueState> {
        let intermediate = self
            .get_users_data(&[username.clone()])
            .remove(username)
            .unwrap_or_default();
        let out = Self::find_appropriate_item(intermediate, flag);
        #[cfg(feature = "runtime_metrics")]
        if out.is_some() {
            self.num_reads.fetch_add(1, Ordering::Relaxed);
        }
        out
    }

    /// Retrieve the batch of specified users user_state's based on the filtering flag provided
    ///
    /// Note: This is a FULL SCAN operation of the entire transaction log
    pub fn get_users_states(
        &self,
        usernames: &[crate::AkdLabel],
        flag: ValueStateRetrievalFlag,
    ) -> HashMap<crate::AkdLabel, ValueState> {
        let mut result_map = HashMap::new();
        let intermediate = self.get_users_data(usernames);

        for (key, value_list) in intermediate.into_iter() {
            if let Some(found) = Self::find_appropriate_item(value_list, flag) {
                result_map.insert(key, found);
            }
        }
        #[cfg(feature = "runtime_metrics")]
        {
            self.num_reads.fetch_add(1, Ordering::Relaxed);
        }
        result_map
    }

    /// Find the appropriate item of the cached value states for a given user. This assumes that the incoming vector
    /// is already sorted in ascending epoch order
    fn find_appropriate_item(
        intermediate: Vec<ValueState>,
        flag: ValueStateRetrievalFlag,
    ) -> Option<ValueState> {
        match flag {
            ValueStateRetrievalFlag::SpecificVersion(version) => intermediate
                .into_iter()
                .find(|item| item.version == version),
            ValueStateRetrievalFlag::SpecificEpoch(epoch) => {
                intermediate.into_iter().find(|item| item.epoch == epoch)
            }
            ValueStateRetrievalFlag::LeqEpoch(epoch) => intermediate
                .into_iter()
                .rev()
                .find(|item| item.epoch <= epoch),
            ValueStateRetrievalFlag::MaxEpoch => intermediate.into_iter().last(),
            ValueStateRetrievalFlag::MinEpoch => intermediate.into_iter().next(),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::append_only_zks::*;
    use crate::node_label::*;
    use crate::storage::types::*;
    use crate::tree_node::*;
    use crate::utils::byte_arr_from_u64;
    use crate::{AkdLabel, AkdValue, AzksValue};
    use rand::{rngs::StdRng, seq::SliceRandom, SeedableRng};

    #[test]
    fn test_commit_order() -> Result<(), StorageError> {
        let azks = DbRecord::Azks(Azks {
            num_nodes: 0,
            latest_epoch: 0,
        });
        let node1 = DbRecord::TreeNode(TreeNodeWithPreviousValue::from_tree_node(TreeNode {
            label: NodeLabel::new(byte_arr_from_u64(0), 0),
            last_epoch: 1,
            min_descendant_epoch: 1,
            parent: NodeLabel::new(byte_arr_from_u64(0), 0),
            node_type: TreeNodeType::Root,
            left_child: None,
            right_child: None,
            hash: AzksValue(crate::hash::EMPTY_DIGEST),
        }));
        let node2 = DbRecord::TreeNode(TreeNodeWithPreviousValue::from_tree_node(TreeNode {
            label: NodeLabel::new(byte_arr_from_u64(1), 1),
            last_epoch: 1,
            min_descendant_epoch: 1,
            parent: NodeLabel::new(byte_arr_from_u64(0), 0),
            node_type: TreeNodeType::Leaf,
            left_child: None,
            right_child: None,
            hash: AzksValue(crate::hash::EMPTY_DIGEST),
        }));
        let value1 = DbRecord::ValueState(ValueState {
            username: AkdLabel::from("test"),
            epoch: 1,
            label: NodeLabel::new(byte_arr_from_u64(1), 1),
            version: 1,
            value: AkdValue::from("abc123"),
        });
        let value2 = DbRecord::ValueState(ValueState {
            username: AkdLabel::from("test"),
            epoch: 2,
            label: NodeLabel::new(byte_arr_from_u64(1), 1),
            version: 2,
            value: AkdValue::from("abc1234"),
        });

        let records = vec![azks, node1, node2, value1, value2];
        let mut rng = StdRng::seed_from_u64(42);

        for _ in 1..10 {
            let txn = Transaction::new();
            txn.begin_transaction();

            // set values in a random order
            let mut shuffled = records.clone();
            shuffled.shuffle(&mut rng);
            for record in shuffled {
                txn.set(&record);
            }

            // ensure that committed records are in ascending priority
            let mut running_priority = 0;
            for record in txn.commit_transaction()? {
                let priority = record.transaction_priority();
                #[allow(clippy::comparison_chain)]
                if priority > running_priority {
                    running_priority = priority;
                } else if priority < running_priority {
                    panic!("Transaction did not obey record priority when committing");
                }
            }
        }

        Ok(())
    }
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed licenses.

//! Various storage and representation related types

use akd_core::AzksValue;

use crate::storage::Storable;
use crate::tree_node::{TreeNode, TreeNodeType, TreeNodeWithPreviousValue};
use crate::{AkdLabel, AkdValue};
use crate::{Azks, NodeLabel};
use std::convert::TryInto;

/// Various elements that can be stored
#[derive(PartialEq, Eq, Debug, Hash, Clone, Copy)]
pub enum StorageType {
    /// Azks
    Azks = 1,
    /// TreeNode
    TreeNode = 2,
    /// EOZ: HistoryNodeState = 3 was removed from here.
    /// Better to keep ValueState = 4 as is?
    /// ValueState
    ValueState = 4,
}

/// State for a value at a given version for that key
#[derive(Debug, Clone, Hash, Eq, PartialEq)]
#[cfg_attr(
    feature = "serde_serialization",
    derive(serde::Deserialize, serde::Serialize)
)]
pub struct ValueStateKey(pub Vec<u8>, pub u64);

/// The state of the value for a given key, starting at a particular epoch.
#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
#[cfg_attr(
    feature = "serde_serialization",
    derive(serde::Deserialize, serde::Serialize)
)]
#[cfg_attr(feature = "serde_serialization", serde(bound = ""))]
pub struct ValueState {
    /// The plaintext value of the user information in the directory
    pub value: AkdValue, // The actual value
    /// The version of the user's value-state
    pub version: u64,
    /// The Node Label
    pub label: NodeLabel,
    /// The epoch this value state was published in
    pub epoch: u64,
    /// The username associated to this value state (username + epoch is the record key)
    pub username: AkdLabel,
}

impl akd_core::SizeOf for ValueState {
    fn size_of(&self) -> usize {
        self.value.size_of()
            + std::mem::size_of::<u64>()
            + self.label.size_of()
            + std::mem::size_of::<u64>()
            + self.username.size_of()
    }
}

impl crate::storage::Storable for ValueState {
    type StorageKey = ValueStateKey;

    fn data_type() -> StorageType {
        StorageType::ValueState
    }

    fn get_id(&self) -> ValueStateKey {
        ValueStateKey(self.username.to_vec(), self.epoch)
    }

    fn get_full_binary_key_id(key: &ValueStateKey) -> Vec<u8> {
        let mut result = vec![StorageType::ValueState as u8];
        result.extend_from_slice(&key.1.to_be_bytes());
        result.extend_from_slice(&key.0);

        result
    }

    fn key_from_full_binary(bin: &[u8]) -> Result<ValueStateKey, String> {
        if bin.len() < 10 {
            return Err("Not enough bytes to form a proper key".to_string());
        }

        if bin[0] != StorageType::ValueState as u8 {
            return Err("Not a value state key".to_string());
        }

        let epoch_bytes: [u8; 8] = bin[1..=8].try_into().expect("Slice with incorrect length");
        let epoch = u64::from_be_bytes(epoch_bytes);
        Ok(ValueStateKey(bin[9..].to_vec(), epoch))
    }
}

impl ValueState {
    pub(crate) fn new(
        username: AkdLabel,
        plaintext_val: AkdValue,
        version: u64,
        label: NodeLabel,
        epoch: u64,
    ) -> Self {
        ValueState {
            value: plaintext_val,
            version,
            label,
            epoch,
            username,
        }
    }
}

/// Data associated with a given key. That is all the states at the various epochs
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
#[cfg_attr(
    feature = "serde_serialization",
    derive(serde::Deserialize, serde::Serialize)
)]
#[cfg_attr(feature = "serde_serialization", serde(bound = ""))]
pub struct KeyData {
    /// The vector of states of key data for a given AkdLabel
    pub states: Vec<ValueState>,
}

/// Used to retrieve a value's state, for a given key
#[derive(std::fmt::Debug, Clone, Copy)]
pub enum ValueStateRetrievalFlag {
    /// Specific version
    SpecificVersion(u64),
    /// State at particular ep
    SpecificEpoch(u64),
    /// State at epoch less than equal to given ep
    LeqEpoch(u64),
    /// State at the latest epoch
    MaxEpoch,
    /// State at the earliest epoch
    MinEpoch,
}

// == New Data Retrieval Logic == //

/// This needs to be PUBLIC public, since anyone implementing a data-layer will need
/// to be able to access this and all the internal types
#[derive(Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
#[cfg_attr(
    feature = "serde_serialization",
    derive(serde::Deserialize, serde::Serialize)
)]
#[allow(clippy::large_enum_variant)]
pub enum DbRecord {
    /// An Azks
    Azks(Azks),
    /// A TreeNode
    TreeNode(TreeNodeWithPreviousValue),
    /// The state of the value for a particular key.
    ValueState(ValueState),
}

impl akd_core::SizeOf for DbRecord {
    fn size_of(&self) -> usize {
        match &self {
            DbRecord::Azks(azks) => azks.size_of(),
            DbRecord::TreeNode(node) => node.size_of(),
            DbRecord::ValueState(state) => state.size_of(),
        }
    }
}

impl Clone for DbRecord {
    fn clone(&self) -> Self {
        match &self {
            DbRecord::Azks(azks) => DbRecord::Azks(azks.clone()),
            DbRecord::TreeNode(node) => DbRecord::TreeNode(node.clone()),
            DbRecord::ValueState(state) => DbRecord::ValueState(state.clone()),
        }
    }
}

impl DbRecord {
    /// Compte a serialized id from the record's fields. This id is useful to use as key
    /// in key-value stores.
    pub fn get_full_binary_id(&self) -> Vec<u8> {
        match &self {
            DbRecord::Azks(azks) => azks.get_full_binary_id(),
            DbRecord::TreeNode(node) => node.get_full_binary_id(),
            DbRecord::ValueState(state) => state.get_full_binary_id(),
        }
    }

    /// Returns the priority in which a record type in a transaction should be committed to storage.
    /// A smaller value indicates higher priority in being written first.
    /// An Azks record should always be updated last, so that any concurrent storage readers will
    /// not see an increase in the current epoch until every other record for the new epoch has
    /// been written to storage.
    pub(crate) fn transaction_priority(&self) -> u8 {
        match &self {
            DbRecord::Azks(_) => 2,
            _ => 1,
        }
    }

    /* Data Layer Builders */

    /// Build an azks instance from the properties
    pub fn build_azks(latest_epoch: u64, num_nodes: u64) -> Azks {
        Azks {
            latest_epoch,
            num_nodes,
        }
    }

    #[allow(clippy::too_many_arguments)]
    /// Build a history tree node from the properties
    pub fn build_tree_node_with_previous_value(
        label_val: [u8; 32],
        label_len: u32,
        last_epoch: u64,
        least_descendant_ep: u64,
        parent_label_val: [u8; 32],
        parent_label_len: u32,
        node_type: u8,
        left_child: Option<NodeLabel>,
        right_child: Option<NodeLabel>,
        value: crate::Digest,
        p_last_epoch: Option<u64>,
        p_least_descendant_ep: Option<u64>,
        p_parent_label_val: Option<[u8; 32]>,
        p_parent_label_len: Option<u32>,
        p_node_type: Option<u8>,
        p_left_child: Option<NodeLabel>,
        p_right_child: Option<NodeLabel>,
        p_value: Option<crate::Digest>,
    ) -> TreeNodeWithPreviousValue {
        let label = NodeLabel::new(label_val, label_len);
        let p_node = match (
            p_last_epoch,
            p_least_descendant_ep,
            p_parent_label_val,
            p_parent_label_len,
            p_node_type,
            p_value,
        ) {
            (Some(a), Some(b), Some(c), Some(d), Some(e), Some(f)) => Some(TreeNode {
                label,
                last_epoch: a,
                min_descendant_epoch: b,
                parent: NodeLabel::new(c, d),
                node_type: TreeNodeType::from_u8(e),
                left_child: p_left_child,
                right_child: p_right_child,
                hash: AzksValue(f),
            }),
            _ => None,
        };
        TreeNodeWithPreviousValue {
            label,
            latest_node: TreeNode {
                label,
                last_epoch,
                min_descendant_epoch: least_descendant_ep,
                parent: NodeLabel::new(parent_label_val, parent_label_len),
                node_type: TreeNodeType::from_u8(node_type),
                left_child,
                right_child,
                hash: AzksValue(value),
            },
            previous_node: p_node,
        }
    }

    /// Build a user state from the properties
    pub fn build_user_state(
        username: Vec<u8>,
        plaintext_val: Vec<u8>,
        version: u64,
        label_len: u32,
        label_val: [u8; 32],
        epoch: u64,
    ) -> ValueState {
        ValueState {
            value: AkdValue(plaintext_val),
            version,
            label: NodeLabel::new(label_val, label_len),
            epoch,
            username: AkdLabel(username),
        }
    }
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed licenses.

//! This module implements a higher-parallelism, async temporary cache for database
//! objects

use super::{CachedItem, DEFAULT_CACHE_CLEAN_FREQUENCY_MS, DEFAULT_ITEM_LIFETIME_MS};
use crate::storage::DbRecord;
use crate::storage::Storable;
use akd_core::SizeOf;
use dashmap::DashMap;
#[cfg(not(feature = "runtime_metrics"))]
use log::debug;
use log::info;
#[cfg(feature = "runtime_metrics")]
use log::{debug, error, warn};

#[cfg(feature = "runtime_metrics")]
use std::sync::atomic::AtomicU64;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::RwLock;

/// Implements a basic cache with timing information which automatically flushes
/// expired entries and removes them
#[derive(Clone)]
pub struct TimedCache {
    azks: Arc<RwLock<Option<DbRecord>>>,
    map: Arc<DashMap<Vec<u8>, CachedItem>>,
    last_clean: Arc<RwLock<Instant>>,
    can_clean: Arc<AtomicBool>,
    item_lifetime: Duration,
    memory_limit_bytes: Option<usize>,
    clean_frequency: Duration,

    #[cfg(feature = "runtime_metrics")]
    hit_count: Arc<AtomicU64>,
}

impl TimedCache {
    /// Log cache access metrics along with size information
    pub fn log_metrics(&self, _level: log::Level) {
        #[cfg(feature = "runtime_metrics")]
        {
            let hit_count = self.hit_count.swap(0, Ordering::Relaxed);
            let cache_size = self.map.len();

            let msg = format!("Cache hit since last: {hit_count}, cached size: {cache_size} items");
            match _level {
                log::Level::Trace => println!("{msg}"),
                log::Level::Debug => debug!("{}", msg),
                log::Level::Info => info!("{}", msg),
                log::Level::Warn => warn!("{}", msg),
                _ => error!("{}", msg),
            }
        }
    }
}

impl TimedCache {
    async fn clean(&self) {
        if !self.can_clean.load(Ordering::Relaxed) {
            // cleaning is disabled
            return;
        }

        let do_clean = {
            // we need the {} brackets in order to release the read lock, since we _may_ acquire a write lock shortly later
            *(self.last_clean.read().await) + self.clean_frequency < Instant::now()
        };
        if do_clean {
            let mut last_clean_write = self.last_clean.write().await;

            let now = Instant::now();
            if let Some(memory_limit_bytes) = self.memory_limit_bytes {
                let mut retained_size = 0;
                let mut num_retained = 0u32;
                let mut num_removed = 0u32;
                self.map.retain(|k, v| {
                    if v.expiration >= now {
                        retained_size += k.len() + v.size_of();
                        num_retained += 1;
                        true
                    } else {
                        num_removed += 1;
                        false
                    }
                });

                info!("Removed {} expired elements from the cache", num_removed);
                debug!("Retained cache size is {} bytes", retained_size);

                if retained_size > memory_limit_bytes {
                    info!("Retained cache size has exceeded the predefined limit, cleaning old entries");
                    // calculate the percentage we'd need to trim off to get to 100% utilization and take another 5%
                    let percent_clean =
                        0.05 + 1.0 - (memory_limit_bytes as f64) / (retained_size as f64);
                    // convert that to the number of items to delete based on the size of the dictionary
                    let num_clean = ((num_retained as f64) * percent_clean).ceil() as usize;
                    // sort the dict based on the oldest entries
                    let mut keys_and_expiration = self
                        .map
                        .iter()
                        .map(|kv| (kv.key().clone(), kv.value().expiration))
                        .collect::<Vec<_>>();
                    keys_and_expiration.sort_by(|(_, a), (_, b)| a.cmp(b));
                    // take `num_clean` old entries and remove them
                    for key in keys_and_expiration
                        .into_iter()
                        .take(num_clean)
                        .map(|(k, _)| k)
                    {
                        self.map.remove(&key);
                    }

                    debug!("END cache memory pressure clean")
                }
            } else {
                // memory pressure analysis is disabled, simply utilize timed cache cleaning
                self.map.retain(|_, v| v.expiration >= now);
            }

            // update last clean time
            *last_clean_write = Instant::now();
        }
    }

    /// Create a new timed cache instance. You can supply an optional item lifetime parameter
    /// or take the default (30s) and an optional memory-pressure limit, where the cache will be
    /// cleaned if too much memory is being utilized
    pub fn new(
        o_lifetime: Option<Duration>,
        o_memory_limit_bytes: Option<usize>,
        o_clean_frequency: Option<Duration>,
    ) -> Self {
        let lifetime = match o_lifetime {
            Some(life) if life > Duration::from_millis(1) => life,
            _ => Duration::from_millis(DEFAULT_ITEM_LIFETIME_MS),
        };
        let clean_frequency = match o_clean_frequency {
            Some(frequency) if frequency > Duration::from_millis(1) => frequency,
            _ => Duration::from_millis(DEFAULT_CACHE_CLEAN_FREQUENCY_MS),
        };
        Self {
            azks: Arc::new(RwLock::new(None)),
            map: Arc::new(DashMap::new()),
            last_clean: Arc::new(RwLock::new(Instant::now())),
            can_clean: Arc::new(AtomicBool::new(true)),
            item_lifetime: lifetime,
            memory_limit_bytes: o_memory_limit_bytes,
            clean_frequency,

            #[cfg(feature = "runtime_metrics")]
            hit_count: Arc::new(AtomicU64::new(0u64)),
        }
    }

    /// Perform a hit-test of the cache for a given key. If successful, Some(record) will be returned
    pub async fn hit_test<St: Storable>(&self, key: &St::StorageKey) -> Option<DbRecord> {
        self.clean().await;

        let full_key = St::get_full_binary_key_id(key);

        // special case for AZKS
        if full_key
            == crate::append_only_zks::Azks::get_full_binary_key_id(
                &crate::append_only_zks::DEFAULT_AZKS_KEY,
            )
        {
            // someone's requesting the AZKS object, return it from the special "cache" storage
            let record = self.azks.read().await.clone();

            #[cfg(feature = "runtime_metrics")]
            self.hit_count.fetch_add(1, Ordering::Relaxed);

            // AZKS objects cannot expire, they need to be manually flushed, so we don't need
            // to check the expiration as below
            return record;
        }

        if let Some(result) = self.map.get(&full_key) {
            #[cfg(feature = "runtime_metrics")]
            self.hit_count.fetch_add(1, Ordering::Relaxed);

            let ignore_clean = !self.can_clean.load(Ordering::Relaxed);
            // if we've disabled cache cleaning, we're in the middle
            // of an in-memory transaction and should ignore expiration
            // of cache items until this flag is disabled again
            if ignore_clean || result.expiration > Instant::now() {
                return Some(result.data.clone());
            }
        }

        None
    }

    /// Put an item into the cache
    pub async fn put(&self, record: &DbRecord) {
        self.clean().await;

        let key = record.get_full_binary_id();

        // special case for AZKS
        if let DbRecord::Azks(azks_ref) = &record {
            let mut guard = self.azks.write().await;
            *guard = Some(DbRecord::Azks(azks_ref.clone()));
        } else {
            let item = CachedItem {
                expiration: Instant::now() + self.item_lifetime,
                data: record.clone(),
            };
            self.map.insert(key, item);
        }
    }

    /// Put a batch of items into the cache, utilizing a single write lock
    pub async fn batch_put(&self, records: &[DbRecord]) {
        self.clean().await;

        for record in records.iter() {
            if let DbRecord::Azks(azks_ref) = &record {
                let mut azks_guard = self.azks.write().await;
                *azks_guard = Some(DbRecord::Azks(azks_ref.clone()));
            } else {
                let key = record.get_full_binary_id();
                let item = CachedItem {
                    expiration: Instant::now() + self.item_lifetime,
                    data: record.clone(),
                };
                self.map.insert(key, item);
            }
        }
    }

    /// Flush the cache
    pub async fn flush(&self) {
        self.map.clear();
        *(self.azks.write().await) = None;
    }

    /// Retrieve all of the cached items
    pub async fn get_all(&self) -> Vec<DbRecord> {
        self.clean().await;

        let mut items = vec![];
        if let Some(record) = self.azks.read().await.clone() {
            items.push(record);
        }
        for kv in self.map.iter() {
            items.push(kv.value().data.clone());
        }

        items
    }

    /// Disable cache-cleaning (i.e. during a transaction)
    pub fn disable_clean(&self) {
        debug!("Disabling cache cleaning");
        self.can_clean.store(false, Ordering::Relaxed);
    }

    /// Re-enable cache cleaning (i.e. when a transaction is over)
    pub fn enable_clean(&self) {
        debug!("Enabling cache cleaning");
        self.can_clean.store(true, Ordering::Relaxed);
    }
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed licenses.

//! This module handles the caching implementation and testing for a time-based cache
//! which supports memory pressure shedding

use crate::storage::DbRecord;
use std::time::Instant;

#[cfg(test)]
mod tests;

/// items live for 30s by default
pub(crate) const DEFAULT_ITEM_LIFETIME_MS: u64 = 30000;
/// clean the cache every 15s by default
pub(crate) const DEFAULT_CACHE_CLEAN_FREQUENCY_MS: u64 = 15000;

pub(crate) struct CachedItem {
    pub(crate) expiration: Instant,
    pub(crate) data: DbRecord,
}

impl akd_core::SizeOf for CachedItem {
    fn size_of(&self) -> usize {
        // the size of an "Instant" varies based on the underlying implementation, so
        // we assume the largest which is 16 bytes on linux
        16 + self.data.size_of()
    }
}

// -------- sub modules -------- //

pub mod high_parallelism;

// -------- cache exports -------- //

pub use high_parallelism::TimedCache;


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed licenses.

//! Caching tests

use super::*;
use std::time::Duration;

use crate::storage::types::{ValueState, ValueStateKey};
use crate::storage::DbRecord;
use crate::{AkdLabel, AkdValue, NodeLabel};

#[tokio::test]
async fn test_cache_put_and_expires() {
    let cache = TimedCache::new(
        Some(Duration::from_millis(10)),
        None,
        Some(Duration::from_millis(50)),
    );

    let value_state = DbRecord::ValueState(ValueState {
        epoch: 1,
        version: 1,
        label: NodeLabel {
            label_len: 1,
            label_val: [0u8; 32],
        },
        value: AkdValue::from("some value"),
        username: AkdLabel::from("user"),
    });
    let key = ValueStateKey(AkdLabel::from("user").0.to_vec(), 1);
    cache.put(&value_state).await;

    let got = cache.hit_test::<ValueState>(&key).await;
    assert!(got.is_some());
    assert_eq!(Some(value_state), got);

    tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
    let got = cache.hit_test::<ValueState>(&key).await;
    assert_eq!(None, got);
}

#[tokio::test]
async fn test_cache_overwrite() {
    let cache = TimedCache::new(Some(Duration::from_millis(1000)), None, None);

    let value_state = ValueState {
        epoch: 1,
        version: 1,
        label: NodeLabel {
            label_len: 1,
            label_val: [0u8; 32],
        },
        value: AkdValue::from("some value"),
        username: AkdLabel::from("user"),
    };
    let key = ValueStateKey(AkdLabel::from("user").0.to_vec(), 1);

    let value_state_2 = ValueState {
        epoch: 1,
        version: 2,
        label: NodeLabel {
            label_len: 2,
            label_val: [0u8; 32],
        },
        value: AkdValue::from("some value"),
        username: AkdLabel::from("user"),
    };
    cache.put(&DbRecord::ValueState(value_state)).await;
    cache
        .put(&DbRecord::ValueState(value_state_2.clone()))
        .await;

    let got = cache.hit_test::<ValueState>(&key).await;
    assert_eq!(Some(DbRecord::ValueState(value_state_2)), got);
}

#[tokio::test]
async fn test_cache_memory_pressure() {
    let cache = TimedCache::new(
        Some(Duration::from_millis(1000)),
        Some(10),
        Some(Duration::from_millis(50)),
    );

    let value_state = DbRecord::ValueState(ValueState {
        epoch: 1,
        version: 1,
        label: NodeLabel {
            label_len: 1,
            label_val: [0u8; 32],
        },
        value: AkdValue::from("some value"),
        username: AkdLabel::from("user"),
    });
    let key = ValueStateKey(AkdLabel::from("user").0.to_vec(), 1);
    cache.put(&value_state).await;

    // we only do an "automated" clean every 50ms in test, which is when memory pressure is evaluated.
    // 100ms will make sure the clean op will run on the next `hit_test` op
    tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
    // This get should return none, even though the cache expiration time is 1s. This is because
    // we should exceed 10 bytes of storage utilization so the cache should clean the item.
    let got = cache.hit_test::<ValueState>(&key).await;
    assert_eq!(None, got);
}

#[tokio::test]
async fn test_many_memory_pressure() {
    let cache = TimedCache::new(
        Some(Duration::from_millis(1000)),
        Some(1024 * 5),
        Some(Duration::from_millis(50)),
    );

    let value_states = (1..100)
        .map(|i| ValueState {
            epoch: i as u64,
            version: i as u64,
            label: NodeLabel {
                label_len: 1,
                label_val: [0u8; 32],
            },
            value: AkdValue::from("test"),
            username: AkdLabel::from("user"),
        })
        .map(DbRecord::ValueState)
        .collect::<Vec<_>>();

    cache.batch_put(&value_states).await;

    // we only do an "automated" clean every 50ms in test, which is when memory pressure is evaluated.
    // 100ms will make sure the clean op will run on the next `hit_test` op
    tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;

    let all = cache.get_all().await;
    assert!(all.len() < 99);
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed licenses.

//! Storage management module for AKD. A wrapper around the underlying database interaction
//! to manage interactions with the data layer to optimize things like caching and
//! transaction management

use crate::storage::cache::TimedCache;
use crate::storage::transaction::Transaction;
use crate::storage::types::DbRecord;
use crate::storage::types::KeyData;
use crate::storage::types::ValueState;
use crate::storage::types::ValueStateKey;
use crate::storage::Database;
use crate::storage::DbSetState;
use crate::storage::Storable;
use crate::storage::StorageError;
use crate::AkdLabel;
use crate::AkdValue;

use log::debug;
#[cfg(feature = "runtime_metrics")]
use log::{error, info, warn};
use std::collections::HashMap;
use std::collections::HashSet;
use std::sync::atomic::AtomicU64;
#[cfg(feature = "runtime_metrics")]
use std::sync::atomic::Ordering;
use std::sync::Arc;
use std::time::Duration;

use super::types::ValueStateRetrievalFlag;

type Metric = usize;

const METRIC_GET: Metric = 0;
const METRIC_BATCH_GET: Metric = 1;
const METRIC_SET: Metric = 2;
const METRIC_BATCH_SET: Metric = 3;
const METRIC_READ_TIME: Metric = 4;
const METRIC_WRITE_TIME: Metric = 5;
const METRIC_TOMBSTONE: Metric = 6;
const METRIC_GET_USER_STATE: Metric = 7;
const METRIC_GET_USER_DATA: Metric = 8;
const METRIC_GET_USER_STATE_VERSIONS: Metric = 9;

const NUM_METRICS: usize = 10;

#[cfg(test)]
mod tests;

/// Represents the manager of the storage mediums, including caching
/// and transactional operations (creating the transaction, committing it, etc)
pub struct StorageManager<Db: Database> {
    cache: Option<TimedCache>,
    transaction: Transaction,
    /// The underlying database managed by this storage manager
    db: Arc<Db>,

    metrics: [Arc<AtomicU64>; NUM_METRICS],
}

impl<Db: Database> Clone for StorageManager<Db> {
    fn clone(&self) -> Self {
        Self {
            cache: self.cache.clone(),
            transaction: self.transaction.clone(),
            db: self.db.clone(),
            metrics: self.metrics.clone(),
        }
    }
}

unsafe impl<Db: Database> Sync for StorageManager<Db> {}
unsafe impl<Db: Database> Send for StorageManager<Db> {}

impl<Db: Database> StorageManager<Db> {
    /// Create a new storage manager with NO CACHE
    pub fn new_no_cache(db: Db) -> Self {
        Self {
            cache: None,
            transaction: Transaction::new(),
            db: Arc::new(db),
            metrics: [0; NUM_METRICS].map(|_| Arc::new(AtomicU64::new(0))),
        }
    }

    /// Create a new storage manager with a cache utilizing the options provided (or defaults)
    pub fn new(
        db: Db,
        cache_item_lifetime: Option<Duration>,
        cache_limit_bytes: Option<usize>,
        cache_clean_frequency: Option<Duration>,
    ) -> Self {
        Self {
            cache: Some(TimedCache::new(
                cache_item_lifetime,
                cache_limit_bytes,
                cache_clean_frequency,
            )),
            transaction: Transaction::new(),
            db: Arc::new(db),
            metrics: [0; NUM_METRICS].map(|_| Arc::new(AtomicU64::new(0))),
        }
    }

    /// Retrieve a reference to the database implementation
    #[cfg(any(test, feature = "public_tests"))]
    pub fn get_db(&self) -> Arc<Db> {
        self.db.clone()
    }

    /// Returns whether the storage manager has a cache
    pub fn has_cache(&self) -> bool {
        self.cache.is_some()
    }

    /// Log metrics from the storage manager (cache, transaction, and storage hit rates etc)
    pub async fn log_metrics(&self, level: log::Level) {
        if let Some(cache) = &self.cache {
            cache.log_metrics(level)
        }

        self.transaction.log_metrics(level);

        #[cfg(feature = "runtime_metrics")]
        {
            let snapshot = self
                .metrics
                .iter()
                .map(|metric| metric.swap(0, Ordering::Relaxed))
                .collect::<Vec<_>>();

            let msg = format!(
                "
===================================================
============ Database operation counts ============
===================================================
    SET {},
    BATCH SET {},
    GET {},
    BATCH GET {}
    TOMBSTONE {}
    GET USER STATE {}
    GET USER DATA {}
    GET USER STATE VERSIONS {}
===================================================
============ Database operation timing ============
===================================================
    TIME READ {} ms
    TIME WRITE {} ms",
                snapshot[METRIC_SET],
                snapshot[METRIC_BATCH_SET],
                snapshot[METRIC_GET],
                snapshot[METRIC_BATCH_GET],
                snapshot[METRIC_TOMBSTONE],
                snapshot[METRIC_GET_USER_STATE],
                snapshot[METRIC_GET_USER_DATA],
                snapshot[METRIC_GET_USER_STATE_VERSIONS],
                snapshot[METRIC_READ_TIME],
                snapshot[METRIC_WRITE_TIME]
            );

            match level {
                // Currently logs cannot be captured unless they are
                // println!. Normally Level::Trace should use the trace! macro.
                log::Level::Trace => println!("{msg}"),
                log::Level::Debug => debug!("{}", msg),
                log::Level::Info => info!("{}", msg),
                log::Level::Warn => warn!("{}", msg),
                _ => error!("{}", msg),
            }
        }
    }

    /// Start an in-memory transaction of changes
    pub fn begin_transaction(&self) -> bool {
        let started = self.transaction.begin_transaction();

        // disable the cache cleaning since we're in a write transaction
        // and will want to keep cache'd objects for the life of the transaction
        if let Some(cache) = &self.cache {
            cache.disable_clean();
        }

        started
    }

    /// Commit a transaction in the database
    pub async fn commit_transaction(&self) -> Result<u64, StorageError> {
        // this retrieves all the trans operations, and "de-activates" the transaction flag
        let records = self.transaction.commit_transaction()?;
        let num_records = records.len();

        // The transaction is now complete (or reverted) and therefore we can re-enable
        // the cache cleaning status
        if let Some(cache) = &self.cache {
            cache.enable_clean();
        }

        if records.is_empty() {
            // no-op, there's nothing to commit
            return Ok(0);
        }

        let _epoch = match records.last() {
            Some(DbRecord::Azks(azks)) => Ok(azks.latest_epoch),
            other => Err(StorageError::Transaction(format!(
                "The last record in the transaction log is NOT an Azks record {other:?}"
            ))),
        }?;

        // update the cache
        if let Some(cache) = &self.cache {
            cache.batch_put(&records).await;
        }

        // Write to the database
        self.tic_toc(
            METRIC_WRITE_TIME,
            self.db.batch_set(records, DbSetState::TransactionCommit),
        )
        .await?;
        self.increment_metric(METRIC_BATCH_SET);
        Ok(num_records as u64)
    }

    /// Rollback a transaction
    pub fn rollback_transaction(&self) -> Result<(), StorageError> {
        self.transaction.rollback_transaction()?;
        // The transaction is being reverted and therefore we can re-enable
        // the cache cleaning status
        if let Some(cache) = &self.cache {
            cache.enable_clean();
        }
        Ok(())
    }

    /// Retrieve a flag determining if there is a transaction active
    pub fn is_transaction_active(&self) -> bool {
        self.transaction.is_transaction_active()
    }

    /// Disable cache cleaning (if present)
    pub fn disable_cache_cleaning(&self) {
        if let Some(cache) = &self.cache {
            cache.disable_clean();
        }
    }

    /// Enable cache cleaning (if present)
    pub fn enable_cache_cleaning(&self) {
        if let Some(cache) = &self.cache {
            cache.enable_clean();
        }
    }

    /// Store a record in the database
    pub async fn set(&self, record: DbRecord) -> Result<(), StorageError> {
        // we're in a transaction, set the item in the transaction
        if self.is_transaction_active() {
            self.transaction.set(&record);
            return Ok(());
        }

        // update the cache
        if let Some(cache) = &self.cache {
            cache.put(&record).await;
        }

        // write to the database
        self.tic_toc(METRIC_WRITE_TIME, self.db.set(record)).await?;
        self.increment_metric(METRIC_SET);
        Ok(())
    }

    /// Set a batch of records in the database
    pub async fn batch_set(&self, records: Vec<DbRecord>) -> Result<(), StorageError> {
        if records.is_empty() {
            // nothing to do, save the cycles
            return Ok(());
        }

        // we're in a transaction, set the items in the transaction
        if self.is_transaction_active() {
            self.transaction.batch_set(&records);
            return Ok(());
        }

        // update the cache
        if let Some(cache) = &self.cache {
            cache.batch_put(&records).await;
        }

        // Write to the database
        self.tic_toc(
            METRIC_WRITE_TIME,
            self.db.batch_set(records, DbSetState::General),
        )
        .await?;
        self.increment_metric(METRIC_BATCH_SET);
        Ok(())
    }

    /// Retrieve a stored record directly from the data layer, ignoring any caching or transaction processes
    pub async fn get_direct<St: Storable>(
        &self,
        id: &St::StorageKey,
    ) -> Result<DbRecord, StorageError> {
        // cache miss, read direct from db
        let record = self
            .tic_toc(METRIC_READ_TIME, self.db.get::<St>(id))
            .await?;
        self.increment_metric(METRIC_GET);
        Ok(record)
    }

    /// Retrieve from the cache only, not falling through to the data-layer. Check's the transaction
    /// if active
    pub async fn get_from_cache_only<St: Storable>(&self, id: &St::StorageKey) -> Option<DbRecord> {
        // we're in a transaction, meaning the object _might_ be newer and therefore we should try and read if from the transaction
        // log instead of the raw storage layer
        if self.is_transaction_active() {
            if let Some(result) = self.transaction.get::<St>(id) {
                return Some(result);
            }
        }

        // check for a cache hit
        if let Some(cache) = &self.cache {
            if let Some(result) = cache.hit_test::<St>(id).await {
                return Some(result);
            }
        }

        None
    }

    /// Retrieve a stored record from the database
    pub async fn get<St: Storable>(&self, id: &St::StorageKey) -> Result<DbRecord, StorageError> {
        if let Some(result) = self.get_from_cache_only::<St>(id).await {
            return Ok(result);
        }

        // cache miss, read direct from db
        self.increment_metric(METRIC_GET);

        let record = self
            .tic_toc(METRIC_READ_TIME, self.db.get::<St>(id))
            .await?;
        if let Some(cache) = &self.cache {
            // cache the result
            cache.put(&record).await;
        }
        Ok(record)
    }

    /// Retrieve a batch of records by id from the database
    pub async fn batch_get<St: Storable>(
        &self,
        ids: &[St::StorageKey],
    ) -> Result<Vec<DbRecord>, StorageError> {
        let mut records = Vec::new();

        if ids.is_empty() {
            // nothing to retrieve, save the cycles
            return Ok(records);
        }

        let mut key_set: HashSet<St::StorageKey> = ids.iter().cloned().collect();

        let trans_active = self.is_transaction_active();
        // first check the transaction log & cache records
        for id in ids.iter() {
            if trans_active {
                // we're in a transaction, meaning the object _might_ be newer and therefore we should try and read if from the transaction
                // log instead of the raw storage layer
                if let Some(result) = self.transaction.get::<St>(id) {
                    records.push(result);
                    key_set.remove(id);
                    continue;
                }
            }

            // check if item is cached
            if let Some(cache) = &self.cache {
                if let Some(result) = cache.hit_test::<St>(id).await {
                    records.push(result);
                    key_set.remove(id);
                    continue;
                }
            }
        }

        if !key_set.is_empty() {
            // these are items to be retrieved from the backing database (not in pending transaction or in the object cache)
            let keys = key_set.into_iter().collect::<Vec<_>>();
            let mut results = self
                .tic_toc(METRIC_READ_TIME, self.db.batch_get::<St>(&keys))
                .await?;

            // cache the db returned results
            if let Some(cache) = &self.cache {
                cache.batch_put(&results).await;
            }

            records.append(&mut results);
            self.increment_metric(METRIC_BATCH_GET);
        }
        Ok(records)
    }

    /// Flush the caching of objects (if present)
    pub async fn flush_cache(&self) {
        if let Some(cache) = &self.cache {
            cache.flush().await;
        }
    }

    /// Tombstone a set of records adhereing to the caching + transactional
    /// settings of the storage manager
    pub async fn tombstone_value_states(&self, keys: &[ValueStateKey]) -> Result<(), StorageError> {
        if keys.is_empty() {
            return Ok(());
        }

        let data = self.batch_get::<ValueState>(keys).await?;
        let mut new_data = vec![];
        for record in data {
            if let DbRecord::ValueState(value_state) = record {
                new_data.push(DbRecord::ValueState(ValueState {
                    epoch: value_state.epoch,
                    label: value_state.label,
                    value: crate::AkdValue(crate::TOMBSTONE.to_vec()),
                    username: value_state.username,
                    version: value_state.version,
                }));
            }
        }
        if !new_data.is_empty() {
            debug!("Tombstoning {} entries", new_data.len());
            self.batch_set(new_data).await?;
            self.increment_metric(METRIC_TOMBSTONE);
        }

        Ok(())
    }

    /// Retrieve the specified user state object based on the retrieval flag from the database
    pub async fn get_user_state(
        &self,
        username: &AkdLabel,
        flag: ValueStateRetrievalFlag,
    ) -> Result<ValueState, StorageError> {
        let maybe_db_state = match self
            .tic_toc(METRIC_READ_TIME, self.db.get_user_state(username, flag))
            .await
        {
            Err(StorageError::NotFound(_)) => Ok(None),
            Ok(something) => Ok(Some(something)),
            Err(other) => Err(other),
        }?;
        self.increment_metric(METRIC_GET_USER_STATE);

        // in the event we are in a transaction, there may be an updated object in the
        // transactional storage. Therefore we should update the db retrieved value if
        // we can with what's in the transaction log
        if self.is_transaction_active() {
            if let Some(transaction_value) = self.transaction.get_user_state(username, flag) {
                if let Some(db_value) = &maybe_db_state {
                    if let Some(record) = Self::compare_db_and_transaction_records(
                        db_value.epoch,
                        transaction_value,
                        flag,
                    ) {
                        return Ok(record);
                    }
                } else {
                    // no db record, but there is a transaction record so use that
                    return Ok(transaction_value);
                }
            }
        }

        if let Some(state) = maybe_db_state {
            // cache the item for future access
            if let Some(cache) = &self.cache {
                cache.put(&DbRecord::ValueState(state.clone())).await;
            }

            Ok(state)
        } else {
            Err(StorageError::NotFound(format!("ValueState {username:?}")))
        }
    }

    /// Retrieve all values states for a given user
    pub async fn get_user_data(&self, username: &AkdLabel) -> Result<KeyData, StorageError> {
        let maybe_db_data = match self
            .tic_toc(METRIC_READ_TIME, self.db.get_user_data(username))
            .await
        {
            Err(StorageError::NotFound(_)) => Ok(None),
            Ok(something) => Ok(Some(something)),
            Err(other) => Err(other),
        }?;
        self.increment_metric(METRIC_GET_USER_DATA);

        if self.is_transaction_active() {
            // there are transaction-based values in the current transaction, they should override database-retrieved values
            let mut map = maybe_db_data
                .map(|data| {
                    data.states
                        .into_iter()
                        .map(|state| (state.epoch, state))
                        .collect::<HashMap<u64, _>>()
                })
                .unwrap_or_else(HashMap::new);

            let transaction_records = self
                .transaction
                .get_users_data(&[username.clone()])
                .remove(username)
                .unwrap_or_default();
            for transaction_record in transaction_records.into_iter() {
                map.insert(transaction_record.epoch, transaction_record);
            }

            return Ok(KeyData {
                states: map.into_values().collect::<Vec<_>>(),
            });
        }

        if let Some(data) = maybe_db_data {
            Ok(data)
        } else {
            Err(StorageError::NotFound(format!(
                "ValueState records for {username:?}"
            )))
        }
    }

    /// Retrieve the user -> state version mapping in bulk. This is the same as get_user_state in a loop, but with less data retrieved from the storage layer
    pub async fn get_user_state_versions(
        &self,
        usernames: &[AkdLabel],
        flag: ValueStateRetrievalFlag,
    ) -> Result<HashMap<AkdLabel, (u64, AkdValue)>, StorageError> {
        let mut data = self
            .tic_toc(
                METRIC_READ_TIME,
                self.db.get_user_state_versions(usernames, flag),
            )
            .await?;
        self.increment_metric(METRIC_GET_USER_STATE_VERSIONS);

        // in the event we are in a transaction, there may be an updated object in the
        // transactional storage. Therefore we should update the db retrieved value if
        // we can with what's in the transaction log
        if self.is_transaction_active() {
            let transaction_records = self.transaction.get_users_states(usernames, flag);
            for (label, value_state) in transaction_records.into_iter() {
                if let Some((epoch, _)) = data.get(&label) {
                    // there is an existing DB record, check if we should updated it from the transaction log
                    if let Some(updated_record) =
                        Self::compare_db_and_transaction_records(*epoch, value_state, flag)
                    {
                        data.insert(label, (*epoch, updated_record.value));
                    }
                } else {
                    // there is no db-equivalent record, but there IS a record in the transaction log.
                    // Take the transaction log value
                    data.insert(label, (value_state.epoch, value_state.value));
                }
            }
        }

        Ok(data)
    }

    fn compare_db_and_transaction_records(
        state_epoch: u64,
        transaction_value: ValueState,
        flag: ValueStateRetrievalFlag,
    ) -> Option<ValueState> {
        match flag {
            ValueStateRetrievalFlag::SpecificVersion(_) => {
                return Some(transaction_value);
            }
            ValueStateRetrievalFlag::SpecificEpoch(_) => {
                return Some(transaction_value);
            }
            ValueStateRetrievalFlag::LeqEpoch(_) => {
                if transaction_value.epoch >= state_epoch {
                    // the transaction has either the same epoch or an epoch in the future, and therefore should
                    // override the db value
                    return Some(transaction_value);
                }
            }
            ValueStateRetrievalFlag::MaxEpoch => {
                if transaction_value.epoch >= state_epoch {
                    // the transaction has either the same epoch or an epoch in the future, and therefore should
                    // override the db value
                    return Some(transaction_value);
                }
            }
            ValueStateRetrievalFlag::MinEpoch => {
                if transaction_value.epoch <= state_epoch {
                    // the transaction has either the same epoch or an older epoch, and therefore should
                    // override the db value
                    return Some(transaction_value);
                }
            }
        }
        None
    }

    fn increment_metric(&self, _metric: Metric) {
        #[cfg(feature = "runtime_metrics")]
        {
            self.metrics[_metric].fetch_add(1, Ordering::Relaxed);
        }
    }

    async fn tic_toc<T>(&self, _metric: Metric, f: impl std::future::Future<Output = T>) -> T {
        #[cfg(feature = "runtime_metrics")]
        {
            let tic = std::time::Instant::now();
            let out = f.await;
            let delta = std::time::Instant::now().duration_since(tic);

            self.metrics[_metric].fetch_add(delta.as_millis() as u64, Ordering::Relaxed);

            out
        }
        #[cfg(not(feature = "runtime_metrics"))]
        {
            f.await
        }
    }
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed licenses.

//! Tests of the storage manager

use akd_core::hash::EMPTY_DIGEST;

use super::*;
use crate::storage::memory::AsyncInMemoryDatabase;
use crate::storage::{types::*, StorageUtil};
use crate::tree_node::{NodeKey, TreeNodeWithPreviousValue};
use crate::*;

#[tokio::test]
async fn test_storage_manager_transaction() {
    let db = AsyncInMemoryDatabase::new();
    let storage_manager = StorageManager::new_no_cache(db);

    assert!(
        storage_manager.begin_transaction(),
        "Failed to start transaction"
    );

    let mut records = (0..10)
        .map(|i| {
            let label = NodeLabel {
                label_len: i,
                label_val: [i as u8; 32],
            };
            DbRecord::TreeNode(DbRecord::build_tree_node_with_previous_value(
                label.label_val,
                label.label_len,
                0,
                0,
                [0u8; 32],
                0,
                0,
                None,
                None,
                EMPTY_DIGEST,
                None,
                None,
                None,
                None,
                None,
                None,
                None,
                None,
            ))
        })
        .collect::<Vec<_>>();

    records.push(DbRecord::Azks(Azks {
        latest_epoch: 0,
        num_nodes: 0,
    }));

    storage_manager
        .batch_set(records)
        .await
        .expect("Failed to set batch of records");
    // there should be no items in the db, as they should all be in the transaction log
    assert_eq!(
        Ok(0),
        storage_manager
            .db
            .batch_get_all_direct()
            .await
            .map(|items| items.len())
    );
    assert_eq!(11, storage_manager.transaction.count());

    // test a retrieval doesn't go to the database. Since we know the db is empty, it should be retrieved from the transaction log
    let key = NodeKey(NodeLabel {
        label_len: 2,
        label_val: [2u8; 32],
    });
    storage_manager
        .get::<TreeNodeWithPreviousValue>(&key)
        .await
        .expect("Failed to get database record for node label 2");

    let keys = vec![
        key,
        NodeKey(NodeLabel {
            label_len: 3,
            label_val: [3u8; 32],
        }),
    ];
    let got = storage_manager
        .batch_get::<TreeNodeWithPreviousValue>(&keys)
        .await
        .expect("Failed to batch-get");
    assert_eq!(2, got.len());

    storage_manager
        .commit_transaction()
        .await
        .expect("Failed to commit transaction");
    // now the records should be in the database and the transaction log empty
    assert_eq!(
        Ok(11),
        storage_manager
            .db
            .batch_get_all_direct()
            .await
            .map(|items| items.len())
    );
    assert_eq!(0, storage_manager.transaction.count());
}

#[tokio::test]
async fn test_storage_manager_cache_populated_by_batch_set() {
    let db = AsyncInMemoryDatabase::new();

    let storage_manager = StorageManager::new(db, None, None, None);

    let mut records = (0..10)
        .map(|i| {
            let label = NodeLabel {
                label_len: i,
                label_val: [i as u8; 32],
            };
            DbRecord::TreeNode(DbRecord::build_tree_node_with_previous_value(
                label.label_val,
                label.label_len,
                0,
                0,
                [0u8; 32],
                0,
                0,
                None,
                None,
                EMPTY_DIGEST,
                None,
                None,
                None,
                None,
                None,
                None,
                None,
                None,
            ))
        })
        .collect::<Vec<_>>();

    records.push(DbRecord::Azks(Azks {
        latest_epoch: 0,
        num_nodes: 0,
    }));

    // write straight to the db, populating the cache
    storage_manager
        .batch_set(records)
        .await
        .expect("Failed to set batch of records");

    // flush the database
    storage_manager.db.clear();

    // test a retrieval still gets data (from the cache)
    let key = NodeKey(NodeLabel {
        label_len: 2,
        label_val: [2u8; 32],
    });
    storage_manager
        .get::<TreeNodeWithPreviousValue>(&key)
        .await
        .expect("Failed to get database record for node label 2");

    let keys = vec![
        key,
        NodeKey(NodeLabel {
            label_len: 3,
            label_val: [3u8; 32],
        }),
    ];
    let got = storage_manager
        .batch_get::<TreeNodeWithPreviousValue>(&keys)
        .await
        .expect("Failed to batch-get");
    assert_eq!(2, got.len());

    storage_manager.flush_cache().await;

    let got = storage_manager
        .batch_get::<TreeNodeWithPreviousValue>(&keys)
        .await
        .expect("Failed to batch-get");
    assert_eq!(0, got.len());
}

#[tokio::test]
async fn test_storage_manager_cache_populated_by_batch_get() {
    let db = AsyncInMemoryDatabase::new();
    let storage_manager = StorageManager::new(db, None, None, None);

    let mut keys = vec![];
    let mut records = (0..10)
        .map(|i| {
            let label = NodeLabel {
                label_len: i,
                label_val: [i as u8; 32],
            };
            keys.push(NodeKey(label));
            DbRecord::TreeNode(DbRecord::build_tree_node_with_previous_value(
                label.label_val,
                label.label_len,
                0,
                0,
                [0u8; 32],
                0,
                0,
                None,
                None,
                EMPTY_DIGEST,
                None,
                None,
                None,
                None,
                None,
                None,
                None,
                None,
            ))
        })
        .collect::<Vec<_>>();

    records.push(DbRecord::Azks(Azks {
        latest_epoch: 0,
        num_nodes: 0,
    }));

    // write straight to the db
    storage_manager
        .batch_set(records)
        .await
        .expect("Failed to set batch of records");

    let db_arc = storage_manager.get_db();
    // flush the cache by destroying the storage manager
    drop(storage_manager);

    // re-create the storage manager, and run a batch_get of the same data keys to populate the cache
    let storage_manager = StorageManager::new(
        Arc::try_unwrap(db_arc).expect("Failed to grab arc"),
        Some(std::time::Duration::from_secs(1000)),
        None,
        None,
    );

    let _ = storage_manager
        .batch_get::<TreeNodeWithPreviousValue>(&keys)
        .await
        .expect("Failed to get a batch of records");

    // flush the database
    storage_manager.db.clear();

    // test a retrieval still gets data (from the cache)
    let key = NodeKey(NodeLabel {
        label_len: 2,
        label_val: [2u8; 32],
    });
    storage_manager
        .get::<TreeNodeWithPreviousValue>(&key)
        .await
        .expect("Failed to get database record for node label 2");

    let keys = vec![
        key,
        NodeKey(NodeLabel {
            label_len: 3,
            label_val: [3u8; 32],
        }),
    ];
    let got = storage_manager
        .batch_get::<TreeNodeWithPreviousValue>(&keys)
        .await
        .expect("Failed to batch-get");
    assert_eq!(2, got.len());

    storage_manager.flush_cache().await;

    // This should be an empty result
    assert_eq!(
        Ok(vec![]),
        storage_manager
            .batch_get::<TreeNodeWithPreviousValue>(&keys)
            .await
    );
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed licenses.

//! # Overview
//!
//! This crate contains a "lean" client to verify AKD proofs which doesn't depend on any
//! crates other than the native hashing implementations and VRF functionality through
//! the [akd_core] crate. This makes it suitable for embedded applications, e.g. inside
//! limited clients (Android, iPhone, WebAssembly, etc) which may not have a large
//! dependency library they can pull upon.
//!
//! ## Features
//!
//! The features of this library are
//!
//! 1. **default** blake3: Blake3 256-bit hashing
//! 2. sha256: SHA2 256-bit hashing
//! 3. sha512: SHA3 512-bit hashing
//! 4. sha3_256: SHA3 256-bit hashing
//! 5. sha3_512: SHA3 512-bit hashing
//!
//! which dictate which hashing function is used by the verification components. Blake3 256-bit hashing is the default
//! implementation and utilizes the `blake3` crate. Features sha256 and sha512 both utilize SHA2 cryptographic functions
//! from the [`sha2`](https://crates.io/crates/sha2) crate. Lastly sha3_256 and sha3_512 features utilize the
//! [`sha3`](https://crates.io/crates/sha3) crate for their hashing implementations.
//! To utilize a hash implementation other than blake3, you should compile with
//!
//! ```bash
//! //          [disable blake3]      [enable other hash]
//! cargo build --no-default-features --features sha3_256
//! ```
//!

#![warn(missing_docs)]
#![cfg_attr(docsrs, feature(doc_cfg))]
#![cfg_attr(feature = "nostd", no_std)]
extern crate alloc;

// Re-expose the core functionality from akd_core (verifications, types, etc)
pub use akd_core::verify;
pub use akd_core::*;

#[cfg(feature = "protobuf")]
pub use akd_core::proto::*;

#[cfg(feature = "wasm")]
pub mod wasm;
#[cfg(feature = "wasm")]
pub use wasm::lookup_verify;