// 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.

//! Implementation of a auditable key directory

use crate::append_only_zks::{Azks, InsertMode};
use crate::ecvrf::{VRFKeyStorage, VRFPublicKey};
use crate::errors::{AkdError, DirectoryError, StorageError};
use crate::helper_structs::LookupInfo;
use crate::storage::manager::StorageManager;
use crate::storage::types::{DbRecord, ValueState, ValueStateRetrievalFlag};
use crate::storage::Database;
use crate::{
    AkdLabel, AkdValue, AppendOnlyProof, AzksElement, Digest, EpochHash, HistoryProof, LookupProof,
    NonMembershipProof, UpdateProof,
};

use crate::crypto::{compute_fresh_azks_value, get_commitment_nonce, stale_azks_value};
use crate::VersionFreshness;
use log::{error, info};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;

/// The representation of a auditable key directory
pub struct Directory<S: Database, V> {
    storage: StorageManager<S>,
    vrf: V,
    /// The cache lock guarantees that the cache is not
    /// flushed mid-proof generation. We allow multiple proof generations
    /// to occur (RwLock.read() operations can have multiple) but we want
    /// to make sure no generations are underway when a cache flush occurs
    /// (in this case we do utilize the write() lock which can only occur 1
    /// at a time and gates further read() locks being acquired during write()).
    cache_lock: Arc<RwLock<()>>,
}

// Manual implementation of Clone, see: https://github.com/rust-lang/rust/issues/41481
impl<S: Database, V: VRFKeyStorage> Clone for Directory<S, V> {
    fn clone(&self) -> Self {
        Self {
            storage: self.storage.clone(),
            vrf: self.vrf.clone(),
            cache_lock: self.cache_lock.clone(),
        }
    }
}

impl<S, V> Directory<S, V>
where
    S: Database + 'static,
    V: VRFKeyStorage,
{
    /// Creates a new (stateless) instance of a auditable key directory.
    /// Takes as input a pointer to the storage being used for this instance.
    /// The state is stored in the storage.
    pub async fn new(storage: StorageManager<S>, vrf: V) -> Result<Self, AkdError> {
        let azks = Directory::<S, V>::get_azks_from_storage(&storage, false).await;

        if azks.is_err() {
            // generate a new azks if one is not found
            let azks = Azks::new::<_>(&storage).await?;
            // store it
            storage.set(DbRecord::Azks(azks.clone())).await?;
        }

        Ok(Directory {
            storage,
            cache_lock: Arc::new(RwLock::new(())),
            vrf,
        })
    }

    /// Updates the directory to include the updated key-value pairs.
    pub async fn publish(&self, updates: Vec<(AkdLabel, AkdValue)>) -> Result<EpochHash, AkdError> {
        // The guard will be dropped at the end of the publish
        let _guard = self.cache_lock.read().await;

        let mut update_set = Vec::<AzksElement>::new();
        let mut user_data_update_set = Vec::<ValueState>::new();

        let mut current_azks = self.retrieve_azks().await?;
        let current_epoch = current_azks.get_latest_epoch();
        let next_epoch = current_epoch + 1;

        let mut keys: Vec<AkdLabel> = updates
            .iter()
            .map(|(akd_label, _val)| akd_label.clone())
            .collect();
        // sort the keys, as inserting in primary-key order is more efficient for MySQL
        keys.sort();

        // we're only using the maximum "version" of the user's state at the last epoch
        // they were seen in the directory. Therefore we've minimized the call to only
        // return a hashmap of AkdLabel => u64 and not retrieving the other data which is not
        // read (i.e. the actual _data_ payload).
        let all_user_versions_retrieved = self
            .storage
            .get_user_state_versions(&keys, ValueStateRetrievalFlag::LeqEpoch(current_epoch))
            .await?;

        info!(
            "Retrieved {} previous user versions of {} requested",
            all_user_versions_retrieved.len(),
            keys.len()
        );

        let vrf_computations = updates
            .iter()
            .flat_map(
                |(akd_label, akd_value)| match all_user_versions_retrieved.get(akd_label) {
                    None => vec![(
                        akd_label.clone(),
                        VersionFreshness::Fresh,
                        1u64,
                        akd_value.clone(),
                    )],
                    Some((latest_version, existing_akd_value)) => {
                        if existing_akd_value == akd_value {
                            // Skip this because the user is trying to re-publish the same value
                            return vec![];
                        }
                        vec![
                            (
                                akd_label.clone(),
                                VersionFreshness::Stale,
                                *latest_version,
                                akd_value.clone(),
                            ),
                            (
                                akd_label.clone(),
                                VersionFreshness::Fresh,
                                *latest_version + 1,
                                akd_value.clone(),
                            ),
                        ]
                    }
                },
            )
            .collect::<Vec<_>>();

        let vrf_map = self
            .vrf
            .get_node_labels(&vrf_computations)
            .await?
            .into_iter()
            .collect::<HashMap<_, _>>();

        let commitment_key = self.derive_commitment_key().await?;

        for ((akd_label, freshness, version, akd_value), node_label) in vrf_map {
            let azks_value = match freshness {
                VersionFreshness::Stale => stale_azks_value(),
                VersionFreshness::Fresh => {
                    compute_fresh_azks_value(&commitment_key, &node_label, version, &akd_value)
                }
            };
            update_set.push(AzksElement {
                label: node_label,
                value: azks_value,
            });

            if freshness == VersionFreshness::Fresh {
                let latest_state =
                    ValueState::new(akd_label, akd_value, version, node_label, next_epoch);
                user_data_update_set.push(latest_state);
            }
        }

        if update_set.is_empty() {
            info!("After filtering for duplicated user information, there is no publish which is necessary (0 updates)");
            // The AZKS has not been updated/mutated at this point, so we can just return the root hash from before
            let root_hash = current_azks.get_root_hash::<_>(&self.storage).await?;
            return Ok(EpochHash(current_epoch, root_hash));
        }

        if !self.storage.begin_transaction() {
            error!("Transaction is already active");
            return Err(AkdError::Storage(StorageError::Transaction(
                "Transaction is already active".to_string(),
            )));
        }
        info!("Starting inserting new leaves");

        if let Err(err) = current_azks
            .batch_insert_nodes::<_>(&self.storage, update_set, InsertMode::Directory)
            .await
        {
            // If we fail to do the batch-leaf insert, we should rollback the transaction so we can try again cleanly.
            // Only fails if transaction is not currently active.
            let _ = self.storage.rollback_transaction();
            // bubble up the err
            return Err(err);
        }

        // batch all the inserts into a single write to storage (in this case it insert's into the transaction log)
        let mut updates = vec![DbRecord::Azks(current_azks.clone())];
        for update in user_data_update_set.into_iter() {
            updates.push(DbRecord::ValueState(update));
        }
        self.storage.batch_set(updates).await?;

        // Commit the transaction
        info!("Committing transaction");
        match self.storage.commit_transaction().await {
            Ok(num_records) => {
                info!("Transaction committed ({} records)", num_records);
            }
            Err(err) => {
                error!("Failed to commit transaction, rolling back");
                let _ = self.storage.rollback_transaction();
                return Err(AkdError::Storage(err));
            }
        };

        let root_hash = current_azks
            .get_root_hash_safe::<_>(&self.storage, next_epoch)
            .await?;

        Ok(EpochHash(next_epoch, root_hash))
    }

    /// Provides proof for correctness of latest version
    ///
    /// * `akd_label`: The target label to generate a lookup proof for
    ///
    /// Returns [Ok((LookupProof, EpochHash))] upon successful generation for the latest version
    /// of the target label's state. [Err(_)] otherwise
    pub async fn lookup(&self, akd_label: AkdLabel) -> Result<(LookupProof, EpochHash), AkdError> {
        // The guard will be dropped at the end of the proof generation
        let _guard = self.cache_lock.read().await;

        let current_azks = self.retrieve_azks().await?;
        let current_epoch = current_azks.get_latest_epoch();
        let lookup_info = self.get_lookup_info(akd_label, current_epoch).await?;

        let root_hash = EpochHash(
            current_epoch,
            current_azks.get_root_hash(&self.storage).await?,
        );
        let proof = self
            .lookup_with_info(&current_azks, lookup_info, false)
            .await?;
        Ok((proof, root_hash))
    }

    /// Generate a lookup proof with the provided target information
    ///
    /// * `current_azks`: The current [Akzs] element
    /// * `lookup_info`: The information to target in the lookup request. Includes all
    /// necessary information to build the proof
    /// * `skip_preload`: Denotes if we should not preload as part of this optimization. Enabled
    /// from bulk lookup proof generation, as it has its own preloading operation
    ///
    /// Returns [Ok(LookupProof)] if the proof generation succeeded, [Err(_)] otherwise
    async fn lookup_with_info(
        &self,
        current_azks: &Azks,
        lookup_info: LookupInfo,
        skip_preload: bool,
    ) -> Result<LookupProof, AkdError> {
        if !skip_preload {
            // Preload nodes needed for lookup.
            #[cfg(feature = "greedy_lookup_preload")]
            {
                current_azks
                    .greedy_preload_lookup_nodes(&self.storage, lookup_info.clone())
                    .await?;
            }
            #[cfg(not(feature = "greedy_lookup_preload"))]
            {
                current_azks
                    .preload_lookup_nodes(&self.storage, &vec![lookup_info.clone()])
                    .await?;
            }
        }
        let label = &lookup_info.value_state.username;
        let current_version = lookup_info.value_state.version;
        let commitment_key = self.derive_commitment_key().await?;
        let plaintext_value = lookup_info.value_state.value;
        let existence_vrf = self
            .vrf
            .get_label_proof(label, VersionFreshness::Fresh, current_version)
            .await?;
        let commitment_label = self.vrf.get_node_label_from_vrf_proof(existence_vrf).await;
        let lookup_proof = LookupProof {
            epoch: lookup_info.value_state.epoch,
            value: plaintext_value.clone(),
            version: lookup_info.value_state.version,
            existence_vrf_proof: existence_vrf.to_bytes().to_vec(),
            existence_proof: current_azks
                .get_membership_proof(&self.storage, lookup_info.existent_label)
                .await?,
            marker_vrf_proof: self
                .vrf
                .get_label_proof(label, VersionFreshness::Fresh, lookup_info.marker_version)
                .await?
                .to_bytes()
                .to_vec(),
            marker_proof: current_azks
                .get_membership_proof(&self.storage, lookup_info.marker_label)
                .await?,
            freshness_vrf_proof: self
                .vrf
                .get_label_proof(label, VersionFreshness::Stale, current_version)
                .await?
                .to_bytes()
                .to_vec(),
            freshness_proof: current_azks
                .get_non_membership_proof(&self.storage, lookup_info.non_existent_label)
                .await?,
            commitment_nonce: get_commitment_nonce(
                &commitment_key,
                &commitment_label,
                lookup_info.value_state.version,
                &plaintext_value,
            )
            .to_vec(),
        };

        Ok(lookup_proof)
    }

    // TODO(eoz): Call proof generations async
    /// Allows efficient batch lookups by preloading necessary nodes for the lookups.
    pub async fn batch_lookup(
        &self,
        akd_labels: &[AkdLabel],
    ) -> Result<(Vec<LookupProof>, EpochHash), AkdError> {
        // The guard will be dropped at the end of the proof generation
        let _guard = self.cache_lock.read().await;

        let current_azks = self.retrieve_azks().await?;
        let current_epoch = current_azks.get_latest_epoch();

        // Take a union of the labels we will need proofs of for each lookup.
        let mut lookup_infos = Vec::new();
        for akd_label in akd_labels {
            // Save lookup info for later use.
            let lookup_info = self
                .get_lookup_info(akd_label.clone(), current_epoch)
                .await?;
            lookup_infos.push(lookup_info.clone());
        }

        // Load nodes needed using the lookup infos.
        current_azks
            .preload_lookup_nodes(&self.storage, &lookup_infos)
            .await?;

        // Ensure we have got all lookup infos needed.
        assert_eq!(akd_labels.len(), lookup_infos.len());

        let root_hash = EpochHash(
            current_epoch,
            current_azks.get_root_hash(&self.storage).await?,
        );

        let mut lookup_proofs = Vec::new();
        for info in lookup_infos.into_iter() {
            lookup_proofs.push(self.lookup_with_info(&current_azks, info, true).await?);
        }

        Ok((lookup_proofs, root_hash))
    }

    async fn build_lookup_info(&self, latest_st: &ValueState) -> Result<LookupInfo, AkdError> {
        let akd_label = &latest_st.username;
        // Need to account for the case where the latest state is
        // added but the database is in the middle of an update
        let version = latest_st.version;
        let marker_version = 1 << get_marker_version(version);
        let existent_label = self
            .vrf
            .get_node_label(akd_label, VersionFreshness::Fresh, version)
            .await?;
        let marker_label = self
            .vrf
            .get_node_label(akd_label, VersionFreshness::Fresh, marker_version)
            .await?;
        let non_existent_label = self
            .vrf
            .get_node_label(akd_label, VersionFreshness::Stale, version)
            .await?;
        Ok(LookupInfo {
            value_state: latest_st.clone(),
            marker_version,
            existent_label,
            marker_label,
            non_existent_label,
        })
    }

    async fn get_lookup_info(
        &self,
        akd_label: AkdLabel,
        epoch: u64,
    ) -> Result<LookupInfo, AkdError> {
        match self
            .storage
            .get_user_state(&akd_label, ValueStateRetrievalFlag::LeqEpoch(epoch))
            .await
        {
            Err(_) => {
                // Need to throw an error
                match std::str::from_utf8(&akd_label) {
                    Ok(name) => Err(AkdError::Storage(StorageError::NotFound(format!(
                        "User {name} at epoch {epoch}"
                    )))),
                    _ => Err(AkdError::Storage(StorageError::NotFound(format!(
                        "User {akd_label:?} at epoch {epoch}"
                    )))),
                }
            }
            Ok(latest_st) => self.build_lookup_info(&latest_st).await,
        }
    }

    /// Takes in the current state of the server and a label.
    /// If the label is present in the current state,
    /// this function returns all the values ever associated with it,
    /// and the epoch at which each value was first committed to the server state.
    /// It also returns the proof of the latest version being served at all times.
    pub async fn key_history(
        &self,
        akd_label: &AkdLabel,
        params: HistoryParams,
    ) -> Result<(HistoryProof, EpochHash), AkdError> {
        // The guard will be dropped at the end of the proof generation
        let _guard = self.cache_lock.read().await;

        let current_azks = self.retrieve_azks().await?;
        let current_epoch = current_azks.get_latest_epoch();
        let mut user_data = self.storage.get_user_data(akd_label).await?.states;

        // reverse sort from highest epoch to lowest
        user_data.sort_by(|a, b| b.epoch.cmp(&a.epoch));

        // apply filters specified by HistoryParams struct
        user_data = match params {
            HistoryParams::Complete => user_data,
            HistoryParams::MostRecent(n) => user_data.into_iter().take(n).collect::<Vec<_>>(),
            HistoryParams::SinceEpoch(epoch) => {
                user_data = user_data
                    .into_iter()
                    .filter(|val| val.epoch >= epoch)
                    .collect::<Vec<_>>();
                // Ordering should be maintained after filtering, but let's re-sort just in case
                user_data.sort_by(|a, b| b.epoch.cmp(&a.epoch));
                user_data
            }
        };

        if user_data.is_empty() {
            let msg = if let Ok(username_str) = std::str::from_utf8(akd_label) {
                format!("User {username_str}")
            } else {
                format!("User {akd_label:?}")
            };
            return Err(AkdError::Storage(StorageError::NotFound(msg)));
        }

        #[cfg(feature = "preload_history")]
        {
            let mut lookup_infos = vec![];
            for ud in user_data.iter() {
                if let Ok(lo) = self.build_lookup_info(ud).await {
                    lookup_infos.push(lo);
                }
            }
            current_azks
                .preload_lookup_nodes(&self.storage, &lookup_infos)
                .await?;
        }

        let mut update_proofs = Vec::<UpdateProof>::new();
        let mut last_version = 0;
        for user_state in user_data {
            // Ignore states in storage that are ahead of current directory epoch
            if user_state.epoch <= current_epoch {
                let proof = self
                    .create_single_update_proof(akd_label, &user_state)
                    .await?;
                update_proofs.push(proof);
                last_version = if user_state.version > last_version {
                    user_state.version
                } else {
                    last_version
                };
            }
        }
        let next_marker = get_marker_version(last_version) + 1;
        let final_marker = get_marker_version(current_epoch);

        let mut until_marker_vrf_proofs = Vec::<Vec<u8>>::new();
        let mut non_existence_until_marker_proofs = Vec::<NonMembershipProof>::new();

        for ver in last_version + 1..(1 << next_marker) {
            let label_for_ver = self
                .vrf
                .get_node_label(akd_label, VersionFreshness::Fresh, ver)
                .await?;
            let non_existence_of_ver = current_azks
                .get_non_membership_proof(&self.storage, label_for_ver)
                .await?;
            non_existence_until_marker_proofs.push(non_existence_of_ver);
            until_marker_vrf_proofs.push(
                self.vrf
                    .get_label_proof(akd_label, VersionFreshness::Fresh, ver)
                    .await?
                    .to_bytes()
                    .to_vec(),
            );
        }

        let mut future_marker_vrf_proofs = Vec::<Vec<u8>>::new();
        let mut non_existence_of_future_marker_proofs = Vec::<NonMembershipProof>::new();

        for marker_power in next_marker..final_marker + 1 {
            let ver = 1 << marker_power;
            let label_for_ver = self
                .vrf
                .get_node_label(akd_label, VersionFreshness::Fresh, ver)
                .await?;
            let non_existence_of_ver = current_azks
                .get_non_membership_proof(&self.storage, label_for_ver)
                .await?;
            non_existence_of_future_marker_proofs.push(non_existence_of_ver);
            future_marker_vrf_proofs.push(
                self.vrf
                    .get_label_proof(akd_label, VersionFreshness::Fresh, ver)
                    .await?
                    .to_bytes()
                    .to_vec(),
            );
        }

        let root_hash = EpochHash(
            current_epoch,
            current_azks.get_root_hash(&self.storage).await?,
        );

        Ok((
            HistoryProof {
                update_proofs,
                until_marker_vrf_proofs,
                non_existence_until_marker_proofs,
                future_marker_vrf_proofs,
                non_existence_of_future_marker_proofs,
            },
            root_hash,
        ))
    }

    /// Poll for changes in the epoch number of the AZKS struct
    /// stored in the storage layer. If an epoch change is detected,
    /// the object cache (if present) is flushed immediately so
    /// that new objects are retrieved from the storage layer against
    /// the "latest" epoch. There is a "special" flow in the storage layer
    /// to do a storage-layer retrieval which ignores the cache
    pub async fn poll_for_azks_changes(
        &self,
        period: tokio::time::Duration,
        change_detected: Option<tokio::sync::mpsc::Sender<()>>,
    ) -> Result<(), AkdError> {
        // Retrieve the same AZKS that all the other calls see (i.e. the version that could be cached
        // at this point). We'll compare this via an uncached call when a change is notified
        let mut last = Directory::<S, V>::get_azks_from_storage(&self.storage, false).await?;

        loop {
            // loop forever polling for changes
            tokio::time::sleep(period).await;

            let latest = Directory::<S, V>::get_azks_from_storage(&self.storage, true).await?;
            if latest.latest_epoch > last.latest_epoch {
                {
                    // acquire a singleton lock prior to flushing the cache to assert that no
                    // cache accesses are underway (i.e. publish/proof generations/etc)
                    let _guard = self.cache_lock.write().await;
                    // flush the cache in its entirety
                    self.storage.flush_cache().await;
                    // re-fetch the azks to load it into cache so when we release the cache lock
                    // others will see the new AZKS loaded up and ready
                    last = Directory::<S, V>::get_azks_from_storage(&self.storage, false).await?;

                    // notify change occurred
                    if let Some(channel) = &change_detected {
                        channel.send(()).await.map_err(|send_err| {
                            AkdError::Storage(StorageError::Connection(format!(
                                "Tokio MPSC sender failed to publish notification with error {send_err:?}"
                            )))
                        })?;
                    }
                    // drop the guard
                }
            }
        }

        #[allow(unreachable_code)]
        Ok(())
    }

    /// Returns an [AppendOnlyProof] for the leaves inserted into the underlying tree between
    /// the epochs `audit_start_ep` and `audit_end_ep`.
    pub async fn audit(
        &self,
        audit_start_ep: u64,
        audit_end_ep: u64,
    ) -> Result<AppendOnlyProof, AkdError> {
        // The guard will be dropped at the end of the proof generation
        let _guard = self.cache_lock.read().await;

        let current_azks = self.retrieve_azks().await?;
        let current_epoch = current_azks.get_latest_epoch();

        if audit_start_ep >= audit_end_ep {
            Err(AkdError::Directory(DirectoryError::InvalidEpoch(format!(
                "Start epoch {audit_start_ep} is greater than or equal the end epoch {audit_end_ep}"
            ))))
        } else if current_epoch < audit_end_ep {
            Err(AkdError::Directory(DirectoryError::InvalidEpoch(format!(
                "End epoch {audit_end_ep} is greater than the current epoch {current_epoch}"
            ))))
        } else {
            self.storage.disable_cache_cleaning();
            let result = current_azks
                .get_append_only_proof::<_>(&self.storage, audit_start_ep, audit_end_ep)
                .await;
            self.storage.enable_cache_cleaning();
            result
        }
    }

    /// Retrieves the [Azks]
    pub(crate) async fn retrieve_azks(&self) -> Result<Azks, crate::errors::AkdError> {
        Directory::<S, V>::get_azks_from_storage(&self.storage, false).await
    }

    async fn get_azks_from_storage(
        storage: &StorageManager<S>,
        ignore_cache: bool,
    ) -> Result<Azks, crate::errors::AkdError> {
        let got = if ignore_cache {
            storage
                .get_direct::<Azks>(&crate::append_only_zks::DEFAULT_AZKS_KEY)
                .await?
        } else {
            storage
                .get::<Azks>(&crate::append_only_zks::DEFAULT_AZKS_KEY)
                .await?
        };
        match got {
            DbRecord::Azks(azks) => Ok(azks),
            _ => {
                error!("No AZKS can be found. You should re-initialize the directory to create a new one");
                Err(AkdError::Storage(StorageError::NotFound(
                    "AZKS not found".to_string(),
                )))
            }
        }
    }

    /// HELPERS ///

    /// Use this function to retrieve the [VRFPublicKey] for this AKD.
    pub async fn get_public_key(&self) -> Result<VRFPublicKey, AkdError> {
        Ok(self.vrf.get_vrf_public_key().await?)
    }

    async fn create_single_update_proof(
        &self,
        akd_label: &AkdLabel,
        user_state: &ValueState,
    ) -> Result<UpdateProof, AkdError> {
        let epoch = user_state.epoch;
        let value = &user_state.value;
        let version = user_state.version;

        let label_at_ep = self
            .vrf
            .get_node_label(akd_label, VersionFreshness::Fresh, version)
            .await?;

        let current_azks = self.retrieve_azks().await?;
        let existence_vrf = self
            .vrf
            .get_label_proof(akd_label, VersionFreshness::Fresh, version)
            .await?;
        let existence_vrf_proof = existence_vrf.to_bytes().to_vec();
        let existence_label = self.vrf.get_node_label_from_vrf_proof(existence_vrf).await;
        let existence_proof = current_azks
            .get_membership_proof(&self.storage, label_at_ep)
            .await?;
        let mut previous_version_proof = Option::None;
        let mut previous_version_vrf_proof = Option::None;
        if version > 1 {
            let prev_label_at_ep = self
                .vrf
                .get_node_label(akd_label, VersionFreshness::Stale, version - 1)
                .await?;
            previous_version_proof = Option::Some(
                current_azks
                    .get_membership_proof(&self.storage, prev_label_at_ep)
                    .await?,
            );
            previous_version_vrf_proof = Option::Some(
                self.vrf
                    .get_label_proof(akd_label, VersionFreshness::Stale, version - 1)
                    .await?
                    .to_bytes()
                    .to_vec(),
            );
        }

        let commitment_key = self.derive_commitment_key().await?;
        let commitment_nonce =
            get_commitment_nonce(&commitment_key, &existence_label, version, value).to_vec();

        Ok(UpdateProof {
            epoch,
            version,
            value: value.clone(),
            existence_vrf_proof,
            existence_proof,
            previous_version_vrf_proof,
            previous_version_proof,
            commitment_nonce,
        })
    }

    /// Gets the root hash at the current epoch.
    pub async fn get_epoch_hash(&self) -> Result<EpochHash, AkdError> {
        let current_azks = self.retrieve_azks().await?;
        let latest_epoch = current_azks.get_latest_epoch();
        let root_hash = current_azks.get_root_hash(&self.storage).await?;
        Ok(EpochHash(latest_epoch, root_hash))
    }

    // We simply hash the VRF private key to derive the commitment key
    async fn derive_commitment_key(&self) -> Result<Digest, AkdError> {
        let raw_key = self.vrf.retrieve().await?;
        let commitment_key = crate::hash::hash(&raw_key);
        Ok(commitment_key)
    }
}

/// A thin newtype which offers read-only interactivity with a [Directory](Directory).
#[derive(Clone)]
pub struct ReadOnlyDirectory<S, V>(Directory<S, V>)
where
    S: Database + Sync + Send,
    V: VRFKeyStorage;

impl<S, V> ReadOnlyDirectory<S, V>
where
    S: Database + 'static,
    V: VRFKeyStorage,
{
    /// Constructs a new instance of [ReadOnlyDirectory]. In the event that an [Azks](crate::append_only_zks::Azks)
    /// does not exist in the storage, or we're unable to retrieve it from storage, then
    /// a [DirectoryError](crate::errors::DirectoryError) will be returned.
    pub async fn new(storage: StorageManager<S>, vrf: V) -> Result<Self, AkdError> {
        let azks = Directory::<S, V>::get_azks_from_storage(&storage, false).await;

        if azks.is_err() {
            return Err(AkdError::Directory(DirectoryError::ReadOnlyDirectory(
                format!(
                    "Cannot start directory in read-only mode when AZKS is missing, error: {:?}",
                    azks.err().take()
                ),
            )));
        }

        Ok(Self(Directory {
            storage,
            cache_lock: Arc::new(RwLock::new(())),
            vrf,
        }))
    }

    /// Read-only access to [Directory::lookup](Directory::lookup).
    pub async fn lookup(&self, uname: AkdLabel) -> Result<(LookupProof, EpochHash), AkdError> {
        self.0.lookup(uname).await
    }

    /// Read-only access to [Directory::batch_lookup](Directory::batch_lookup).
    pub async fn batch_lookup(
        &self,
        unames: &[AkdLabel],
    ) -> Result<(Vec<LookupProof>, EpochHash), AkdError> {
        self.0.batch_lookup(unames).await
    }

    /// Read-only access to [Directory::key_history](Directory::key_history).
    pub async fn key_history(
        &self,
        uname: &AkdLabel,
        params: HistoryParams,
    ) -> Result<(HistoryProof, EpochHash), AkdError> {
        self.0.key_history(uname, params).await
    }

    /// Read-only access to [Directory::poll_for_azks_changes](Directory::poll_for_azks_changes).
    pub async fn poll_for_azks_changes(
        &self,
        period: tokio::time::Duration,
        change_detected: Option<tokio::sync::mpsc::Sender<()>>,
    ) -> Result<(), AkdError> {
        self.0.poll_for_azks_changes(period, change_detected).await
    }

    /// Read-only access to [Directory::audit](Directory::audit).
    pub async fn audit(
        &self,
        audit_start_ep: u64,
        audit_end_ep: u64,
    ) -> Result<AppendOnlyProof, AkdError> {
        self.0.audit(audit_start_ep, audit_end_ep).await
    }

    /// Read-only access to [Directory::get_epoch_hash].
    pub async fn get_epoch_hash(&self) -> Result<EpochHash, AkdError> {
        self.0.get_epoch_hash().await
    }

    /// Read-only access to [Directory::get_public_key](Directory::get_public_key).
    pub async fn get_public_key(&self) -> Result<VRFPublicKey, AkdError> {
        self.0.get_public_key().await
    }
}

/// The parameters that dictate how much of the history proof to return to the consumer
/// (either a complete history, or some limited form).
#[derive(Copy, Clone)]
pub enum HistoryParams {
    /// Returns a complete history for a label
    Complete,
    /// Returns up to the most recent N updates for a label
    MostRecent(usize),
    /// Returns all updates since a specified epoch (inclusive)
    SinceEpoch(u64),
}

impl Default for HistoryParams {
    /// By default, we return a complete history
    fn default() -> Self {
        Self::Complete
    }
}

/// Helpers

pub(crate) fn get_marker_version(version: u64) -> u64 {
    (64 - version.leading_zeros() - 1).into()
}

/// Helpers for testing

/// This enum is meant to insert corruptions into a malicious publish function.
#[derive(Debug, Clone)]
pub enum PublishCorruption {
    /// Indicates to the malicious publish function to not mark a stale version
    UnmarkedStaleVersion(AkdLabel),
    /// Indicates to the malicious publish to mark a certain version for a username as stale.
    MarkVersionStale(AkdLabel, u64),
}

#[cfg(test)]
impl<S: Database + 'static, V: VRFKeyStorage> Directory<S, V> {
    /// Updates the directory to include the updated key-value pairs with possible issues.
    pub(crate) async fn publish_malicious_update(
        &self,
        updates: Vec<(AkdLabel, AkdValue)>,
        corruption: PublishCorruption,
    ) -> Result<EpochHash, AkdError> {
        // The guard will be dropped at the end of the publish
        let _guard = self.cache_lock.read().await;

        let mut update_set = Vec::<AzksElement>::new();

        if let PublishCorruption::MarkVersionStale(ref akd_label, version_number) = corruption {
            // In the malicious case, sometimes the server may not mark the old version stale immediately.
            // If this is the case, it may want to do this marking at a later time.
            let stale_label = self
                .vrf
                .get_node_label(akd_label, VersionFreshness::Stale, version_number)
                .await?;
            let stale_value_to_add = stale_azks_value();
            update_set.push(AzksElement {
                label: stale_label,
                value: stale_value_to_add,
            })
        };

        let mut user_data_update_set = Vec::<ValueState>::new();

        let mut current_azks = self.retrieve_azks().await?;
        let current_epoch = current_azks.get_latest_epoch();
        let next_epoch = current_epoch + 1;

        let mut keys: Vec<AkdLabel> = updates
            .iter()
            .map(|(akd_label, _val)| akd_label.clone())
            .collect();
        // sort the keys, as inserting in primary-key order is more efficient for MySQL
        keys.sort();

        // we're only using the maximum "version" of the user's state at the last epoch
        // they were seen in the directory. Therefore we've minimized the call to only
        // return a hashmap of AkdLabel => u64 and not retrieving the other data which is not
        // read (i.e. the actual _data_ payload).
        let all_user_versions_retrieved = self
            .storage
            .get_user_state_versions(&keys, ValueStateRetrievalFlag::LeqEpoch(current_epoch))
            .await?;

        info!(
            "Retrieved {} previous user versions of {} requested",
            all_user_versions_retrieved.len(),
            keys.len()
        );

        let commitment_key = self.derive_commitment_key().await?;

        for (akd_label, val) in updates {
            match all_user_versions_retrieved.get(&akd_label) {
                None => {
                    // no data found for the user
                    let latest_version = 1;
                    let label = self
                        .vrf
                        .get_node_label(&akd_label, VersionFreshness::Fresh, latest_version)
                        .await?;

                    let value_to_add =
                        compute_fresh_azks_value(&commitment_key, &label, latest_version, &val);
                    update_set.push(AzksElement {
                        label,
                        value: value_to_add,
                    });
                    let latest_state =
                        ValueState::new(akd_label, val, latest_version, label, next_epoch);
                    user_data_update_set.push(latest_state);
                }
                Some((_, previous_value)) if val == *previous_value => {
                    // skip this version because the user is trying to re-publish the already most recent value
                    // Issue #197: https://github.com/novifinancial/akd/issues/197
                }
                Some((previous_version, _)) => {
                    // Data found for the given user
                    let latest_version = *previous_version + 1;
                    let stale_label = self
                        .vrf
                        .get_node_label(&akd_label, VersionFreshness::Stale, *previous_version)
                        .await?;
                    let fresh_label = self
                        .vrf
                        .get_node_label(&akd_label, VersionFreshness::Fresh, latest_version)
                        .await?;
                    let stale_value_to_add = stale_azks_value();
                    let fresh_value_to_add = compute_fresh_azks_value(
                        &commitment_key,
                        &fresh_label,
                        latest_version,
                        &val,
                    );
                    match &corruption {
                        // Some malicious server might not want to mark an old and compromised key as stale.
                        // Thus, you only push the key if either the corruption is for some other username,
                        // or the corruption is not of the type that asks you to delay marking a stale value correctly.
                        PublishCorruption::UnmarkedStaleVersion(target_akd_label) => {
                            if *target_akd_label != akd_label {
                                update_set.push(AzksElement {
                                    label: stale_label,
                                    value: stale_value_to_add,
                                })
                            }
                        }
                        _ => update_set.push(AzksElement {
                            label: stale_label,
                            value: stale_value_to_add,
                        }),
                    };

                    update_set.push(AzksElement {
                        label: fresh_label,
                        value: fresh_value_to_add,
                    });
                    let new_state =
                        ValueState::new(akd_label, val, latest_version, fresh_label, next_epoch);
                    user_data_update_set.push(new_state);
                }
            }
        }
        let azks_element_set: Vec<AzksElement> = update_set.to_vec();

        if azks_element_set.is_empty() {
            info!("After filtering for duplicated user information, there is no publish which is necessary (0 updates)");
            // The AZKS has not been updated/mutated at this point, so we can just return the root hash from before
            let root_hash = current_azks.get_root_hash::<_>(&self.storage).await?;
            return Ok(EpochHash(current_epoch, root_hash));
        }

        if let false = self.storage.begin_transaction() {
            error!("Transaction is already active");
            return Err(AkdError::Storage(StorageError::Transaction(
                "Transaction is already active".to_string(),
            )));
        }
        info!("Starting database insertion");

        current_azks
            .batch_insert_nodes::<_>(&self.storage, azks_element_set, InsertMode::Directory)
            .await?;

        // batch all the inserts into a single transactional write to storage
        let mut updates = vec![DbRecord::Azks(current_azks.clone())];
        for update in user_data_update_set.into_iter() {
            updates.push(DbRecord::ValueState(update));
        }
        self.storage.batch_set(updates).await?;

        // now commit the transaction
        if let Err(err) = self.storage.commit_transaction().await {
            // ignore any rollback error(s)
            let _ = self.storage.rollback_transaction();
            return Err(AkdError::Storage(err));
        }

        let root_hash = current_azks
            .get_root_hash_safe::<_>(&self.storage, next_epoch)
            .await?;

        Ok(EpochHash(next_epoch, root_hash))
        // At the moment the tree root is not being written anywhere. Eventually we
        // want to change this to call a write operation to post to a blockchain or some such thing
    }
}


// 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.

//! Errors for various data structure operations.
use core::fmt;

use crate::node_label::NodeLabel;
use crate::Direction;

/// Symbolizes a AkdError, thrown by the akd.
#[cfg_attr(test, derive(PartialEq, Eq))]
#[derive(Debug)]
pub enum AkdError {
    /// Error propagation
    TreeNode(TreeNodeError),
    /// Error propagation
    Directory(DirectoryError),
    /// Error propagation
    AzksErr(AzksError),
    /// Vrf related error
    Vrf(akd_core::ecvrf::VrfError),
    /// Storage layer error thrown
    Storage(StorageError),
    /// Audit verification error thrown
    AuditErr(AuditorError),
    /// Parallelism/concurrency related errors
    Parallelism(ParallelismError),
    /// Test error
    TestErr(String),
}

impl std::error::Error for AkdError {}

impl From<TreeNodeError> for AkdError {
    fn from(error: TreeNodeError) -> Self {
        Self::TreeNode(error)
    }
}

impl From<StorageError> for AkdError {
    fn from(error: StorageError) -> Self {
        Self::Storage(error)
    }
}

impl From<DirectoryError> for AkdError {
    fn from(error: DirectoryError) -> Self {
        Self::Directory(error)
    }
}

impl From<akd_core::ecvrf::VrfError> for AkdError {
    fn from(error: akd_core::ecvrf::VrfError) -> Self {
        Self::Vrf(error)
    }
}

impl From<AzksError> for AkdError {
    fn from(error: AzksError) -> Self {
        Self::AzksErr(error)
    }
}

impl From<AuditorError> for AkdError {
    fn from(error: AuditorError) -> Self {
        Self::AuditErr(error)
    }
}

impl From<ParallelismError> for AkdError {
    fn from(error: ParallelismError) -> Self {
        Self::Parallelism(error)
    }
}

impl From<akd_core::verify::VerificationError> for AkdError {
    fn from(err: akd_core::verify::VerificationError) -> Self {
        Self::Directory(err.into())
    }
}

impl std::fmt::Display for AkdError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> Result<(), std::fmt::Error> {
        match self {
            AkdError::TreeNode(err) => {
                writeln!(f, "AKD Tree Node Error: {err}")
            }
            AkdError::Directory(err) => {
                writeln!(f, "AKD Directory Error: {err}")
            }
            AkdError::AzksErr(err) => {
                writeln!(f, "AKD AZKS Error: {err}")
            }
            AkdError::Vrf(err) => {
                writeln!(f, "AKD VRF Error: {err}")
            }
            AkdError::Storage(err) => {
                writeln!(f, "AKD Storage Error: {err}")
            }
            AkdError::AuditErr(err) => {
                writeln!(f, "AKD Auditor Error {err}")
            }
            AkdError::Parallelism(err) => {
                writeln!(f, "AKD Parallelism Error: {err}")
            }
            AkdError::TestErr(err) => {
                writeln!(f, "{err}")
            }
        }
    }
}

/// Errors thrown by TreeNodes
#[derive(Debug, Eq, PartialEq)]
pub enum TreeNodeError {
    /// At the moment the only supported dirs are 0, 1
    InvalidDirection(Direction),
    /// No direction provided for the node.
    /// Second parameter is the label of the child attempted to be set
    /// -- if there is one, otherwise it is None.
    NoDirection(NodeLabel, Option<NodeLabel>),
    /// The node didn't have a child in the given epoch
    NoChildAtEpoch(u64, Direction),
    /// The next epoch of this node's parent was invalid
    ParentNextEpochInvalid(u64),
    /// The hash of a parent was attempted to be updated, without setting the calling node as a child.
    HashUpdateOrderInconsistent,
    /// The node did not exist at epoch
    NonexistentAtEpoch(NodeLabel, u64),
    /// The state of a node did not exist at a given epoch
    NoStateAtEpoch(NodeLabel, u64),
    /// Failed to deserialize a digest
    DigestDeserializationFailed(String),
}

impl std::error::Error for TreeNodeError {}

impl fmt::Display for TreeNodeError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::InvalidDirection(dir) => {
                write!(
                    f,
                    "AKD is based on a binary tree. No child with a given direction: {dir:?}"
                )
            }
            Self::NoDirection(node_label, child_label) => {
                let mut to_print = format!("no direction provided for the node {node_label:?}");
                // Add child info if given.
                if let Some(child_label) = child_label {
                    let child_str = format!(" and child {child_label:?}");
                    to_print.push_str(&child_str);
                }
                write!(f, "{to_print}")
            }
            Self::NoChildAtEpoch(epoch, direction) => {
                write!(f, "no node in direction {direction:?} at epoch {epoch}")
            }
            Self::ParentNextEpochInvalid(epoch) => {
                write!(f, "Next epoch of parent is invalid, epoch = {epoch}")
            }
            Self::HashUpdateOrderInconsistent => {
                write!(
                    f,
                    "Hash update in parent only allowed after node is inserted"
                )
            }
            Self::NonexistentAtEpoch(label, epoch) => {
                write!(
                    f,
                    "This node, labelled {label:?}, did not exist at epoch {epoch:?}."
                )
            }
            Self::NoStateAtEpoch(label, epoch) => {
                write!(
                    f,
                    "This node, labelled {label:?}, did not exist at epoch {epoch:?}."
                )
            }
            Self::DigestDeserializationFailed(inner_error) => {
                write!(f, "Encountered a serialization error {inner_error}")
            }
        }
    }
}

/// An error thrown by the Azks data structure.
#[cfg_attr(test, derive(PartialEq, Eq))]
#[derive(Debug)]
pub enum AzksError {
    /// Membership proof did not verify
    VerifyMembershipProof(String),
    /// Append-only proof did not verify
    VerifyAppendOnlyProof,
    /// Thrown when a place where an epoch is needed wasn't provided one.
    NoEpochGiven,
}

impl std::error::Error for AzksError {}

impl fmt::Display for AzksError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::VerifyMembershipProof(error_string) => {
                write!(f, "{error_string}")
            }
            Self::VerifyAppendOnlyProof => {
                write!(f, "Append only proof did not verify!")
            }
            Self::NoEpochGiven => {
                write!(f, "An epoch was required but not supplied")
            }
        }
    }
}

/// The errors thrown by various algorithms in [crate::directory::Directory]
#[cfg_attr(test, derive(PartialEq, Eq))]
#[derive(Debug)]
pub enum DirectoryError {
    /// A verification error occurred
    Verification(akd_core::verify::VerificationError),
    /// Tried to perform an operation on an invalid epoch or epoch range
    InvalidEpoch(String),
    /// AZKS not found in read-only directory mode
    ReadOnlyDirectory(String),
}

impl std::error::Error for DirectoryError {}

impl fmt::Display for DirectoryError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Verification(err) => {
                write!(f, "Verification failure {err}")
            }
            Self::InvalidEpoch(err_string) => {
                write!(f, "Invalid epoch {err_string}")
            }
            Self::ReadOnlyDirectory(inner_message) => {
                write!(f, "Directory in read-only mode: {inner_message}")
            }
        }
    }
}

impl From<akd_core::verify::VerificationError> for DirectoryError {
    fn from(err: akd_core::verify::VerificationError) -> Self {
        Self::Verification(err)
    }
}

/// Represents a storage-layer error
#[cfg_attr(any(test, feature = "public_tests"), derive(PartialEq, Eq))]
#[derive(Debug)]
pub enum StorageError {
    /// Data wasn't found in the storage layer
    NotFound(String),
    /// A transaction error
    Transaction(String),
    /// Some kind of storage connection error occurred
    Connection(String),
    /// Some other storage-layer error occurred
    Other(String),
}

impl std::error::Error for StorageError {}

impl fmt::Display for StorageError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            StorageError::Connection(inner) => {
                write!(f, "Storage connection: {inner}")
            }
            StorageError::Transaction(inner) => {
                write!(f, "Transaction: {inner}")
            }
            StorageError::NotFound(inner) => {
                write!(f, "Data not found: {inner}")
            }
            StorageError::Other(inner) => {
                write!(f, "Other storage error: {inner}")
            }
        }
    }
}

/// The errors thrown by various algorithms in [crate::directory::Directory]
#[cfg_attr(test, derive(PartialEq, Eq))]
#[derive(Debug)]
pub enum AuditorError {
    /// A general auditor error
    VerifyAuditProof(String),
}

impl std::error::Error for AuditorError {}

impl fmt::Display for AuditorError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::VerifyAuditProof(err_string) => {
                write!(f, "Failed to verify audit {err_string}")
            }
        }
    }
}

/// The errors thrown by parallel code
#[cfg_attr(test, derive(PartialEq, Eq))]
#[derive(Debug)]
pub enum ParallelismError {
    /// A tokio task join error
    JoinErr(String),
}

impl std::error::Error for ParallelismError {}

impl fmt::Display for ParallelismError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::JoinErr(err_string) => {
                write!(f, "Failed to join tokio task {err_string}")
            }
        }
    }
}


// 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.

//! Helper structs that are used for various data structures,
//! to make it easier to pass arguments around.

use crate::Digest;
use crate::{storage::types::ValueState, NodeLabel};

/// Root hash of the tree and its associated epoch
#[derive(Debug, Clone, Hash, PartialEq, Eq)]
pub struct EpochHash(pub u64, pub Digest);

impl EpochHash {
    /// Get the contained epoch
    pub fn epoch(&self) -> u64 {
        self.0
    }
    /// Get the contained hash
    pub fn hash(&self) -> Digest {
        self.1
    }
}

#[derive(Clone, Debug)]
/// Info needed for a lookup of a user for an epoch
pub struct LookupInfo {
    pub(crate) value_state: ValueState,
    pub(crate) marker_version: u64,
    pub(crate) existent_label: NodeLabel,
    pub(crate) marker_label: NodeLabel,
    pub(crate) non_existent_label: NodeLabel,
}


// 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 contains all the type conversions between internal AKD & message types
//! with the protobuf types
//!
//! Additionally it supports the conversion between the output from the `Directory` to
//! public-storage safe blob types encoded with Protobuf. Download and upload
//! to the blob storage medium is left to the new application crate akd_local_auditor

use crate::Digest;
use protobuf::Message;
use std::convert::{TryFrom, TryInto};

/// Local audit processing errors
#[derive(Debug)]
pub enum LocalAuditorError {
    /// An error parsing the blob name to/from a string
    NameParseError(String),
    /// An error between the lengths of hashes + proofs
    MisMatchedLengths(String),
    /// A conversion error occurred
    ConversionError(akd_core::proto::ConversionError),
}

impl From<akd_core::proto::ConversionError> for LocalAuditorError {
    fn from(err: akd_core::proto::ConversionError) -> Self {
        Self::ConversionError(err)
    }
}

impl From<protobuf::Error> for LocalAuditorError {
    fn from(err: protobuf::Error) -> Self {
        Self::ConversionError(err.into())
    }
}

// ************************ Converters ************************ //

macro_rules! hash_from_ref {
    ($obj:expr) => {
        crate::hash::try_parse_digest($obj)
            .map_err(akd_core::proto::ConversionError::Deserialization)
    };
}

// ************************ Helper Functions ************************ //

const NAME_SEPARATOR: char = '/';

/// Represents the NAME of an audit blob and can be
/// flatted to/from a string
#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Hash, Default, Copy)]
pub struct AuditBlobName {
    /// The epoch this audit proof is related to
    pub epoch: u64,
    /// The previous root hash from `&self.epoch - 1`
    pub previous_hash: Digest,
    /// The current updated root hash
    pub current_hash: Digest,
}

impl std::string::ToString for AuditBlobName {
    fn to_string(&self) -> String {
        let previous_hash = hex::encode(self.previous_hash);
        let current_hash = hex::encode(self.current_hash);
        format!(
            "{}{}{}{}{}",
            self.epoch, NAME_SEPARATOR, previous_hash, NAME_SEPARATOR, current_hash
        )
    }
}

impl TryFrom<&str> for AuditBlobName {
    type Error = LocalAuditorError;

    fn try_from(name: &str) -> Result<Self, Self::Error> {
        let parts = name.split(NAME_SEPARATOR).collect::<Vec<_>>();
        if parts.len() < 3 {
            return Err(LocalAuditorError::NameParseError(
                "Name is malformed, there are not enough components to reconstruct!".to_string(),
            ));
        }
        // PART[0] = EPOCH
        let epoch: u64 = parts[0].parse().map_err(|_| {
            LocalAuditorError::NameParseError(format!("Failed to parse '{}' into an u64", parts[0]))
        })?;

        // PART[1] = PREVIOUS_HASH
        let previous_hash_bytes = hex::decode(parts[1]).map_err(|hex_err| {
            LocalAuditorError::NameParseError(format!(
                "Failed to decode previous hash from hex string: {hex_err}"
            ))
        })?;
        let previous_hash = hash_from_ref!(&previous_hash_bytes)?;

        // PART[2] = CURRENT_HASH
        let current_hash_bytes = hex::decode(parts[2]).map_err(|hex_err| {
            LocalAuditorError::NameParseError(format!(
                "Failed to decode current hash from hex string: {hex_err}"
            ))
        })?;
        let current_hash = hash_from_ref!(&current_hash_bytes)?;

        Ok(AuditBlobName {
            epoch,
            current_hash,
            previous_hash,
        })
    }
}

/// The constructed blobs with naming encoding the
/// blob name = "EPOCH/PREVIOUS_ROOT_HASH/CURRENT_ROOT_HASH"
#[derive(Clone)]
pub struct AuditBlob {
    /// The name of the blob, which can be decomposed into logical components (phash, chash, epoch)
    pub name: AuditBlobName,
    /// The binary data comprising the blob contents
    pub data: Vec<u8>,
}

impl AuditBlob {
    /// Construct a new AuditBlob from the internal structures, which is ready to be written to persistent storage
    pub fn new(
        previous_hash: Digest,
        current_hash: Digest,
        epoch: u64,
        proof: &crate::SingleAppendOnlyProof,
    ) -> Result<AuditBlob, LocalAuditorError> {
        let name = AuditBlobName {
            epoch,
            previous_hash,
            current_hash,
        };
        let proto: akd_core::proto::specs::types::SingleAppendOnlyProof = proof.into();

        Ok(AuditBlob {
            name,
            data: proto.write_to_bytes()?,
        })
    }

    /// Decode a protobuf encoded AuditBlob into it's components (phash, chash, epoch, proof)
    pub fn decode(
        &self,
    ) -> Result<(u64, Digest, Digest, crate::SingleAppendOnlyProof), LocalAuditorError> {
        let proof =
            akd_core::proto::specs::types::SingleAppendOnlyProof::parse_from_bytes(&self.data)?;
        let local_proof: crate::SingleAppendOnlyProof = (&proof).try_into()?;

        Ok((
            self.name.epoch,
            hash_from_ref!(&self.name.previous_hash)?,
            hash_from_ref!(&self.name.current_hash)?,
            local_proof,
        ))
    }
}

/// Convert an append-only proof to "Audit Blobs" which are to be stored in a publicly readable storage medium
/// suitable for public auditing
pub fn generate_audit_blobs(
    hashes: Vec<Digest>,
    proof: crate::AppendOnlyProof,
) -> Result<Vec<AuditBlob>, LocalAuditorError> {
    if proof.epochs.len() + 1 != hashes.len() {
        return Err(LocalAuditorError::MisMatchedLengths(format!(
            "The proof has a different number of epochs than needed for hashes.
            The number of hashes you provide should be one more than the number of epochs!
            Number of epochs = {}, number of hashes = {}",
            proof.epochs.len(),
            hashes.len()
        )));
    }

    if proof.epochs.len() != proof.proofs.len() {
        return Err(LocalAuditorError::MisMatchedLengths(format!(
            "The proof has {} epochs and {} proofs. These should be equal!",
            proof.epochs.len(),
            proof.proofs.len()
        )));
    }

    let mut results = Vec::with_capacity(proof.proofs.len());

    for i in 0..hashes.len() - 1 {
        let previous_hash = hashes[i];
        let current_hash = hashes[i + 1];
        // The epoch provided is the source epoch, i.e. the proof is validating from (T, T+1)
        let epoch = proof.epochs[i];

        let blob = AuditBlob::new(previous_hash, current_hash, epoch, &proof.proofs[i])?;
        results.push(blob);
    }

    Ok(results)
}

#[cfg(test)]
mod tests {
    use super::{AuditBlobName, LocalAuditorError};
    use std::convert::TryInto;

    #[test]
    fn test_audit_proof_naming_conventions() -> Result<(), LocalAuditorError> {
        let expected_name = "54/0101010101010101010101010101010101010101010101010101010101010101/0000000000000000000000000000000000000000000000000000000000000000";

        let blob_name = AuditBlobName {
            current_hash: crate::hash::EMPTY_DIGEST,
            previous_hash: [1u8; crate::hash::DIGEST_BYTES],
            epoch: 54,
        };

        let name = blob_name.to_string();
        assert_ne!(String::new(), name);

        assert_eq!(expected_name.to_string(), blob_name.to_string());

        let blob_name_ref: &str = name.as_ref();
        let decomposed: AuditBlobName = blob_name_ref.try_into()?;
        assert_eq!(blob_name, decomposed);
        Ok(())
    }
}