use cid::multihash::MultihashDigest;
use std::collections::BTreeMap;
use std::result;
use std::sync::Arc;

use ex3_balance_vault_public_types::{
    AcceptedWithdrawalShardingSubReport, BalanceDataIntegrityShardingSubReport,
    CandidBalanceShardingSubReport, RejectedWithdrawalShardingSubReport,
};
use ex3_block_reports::BlockWithShardingReports;
use ex3_blockchain_public_types::{
    CandidConsensusReport, IdCursorReport, NewSnapshotVaultActivatedReport, SnapshotCursor,
    WithdrawalCounterReport,
};
use ex3_core_registry_public_types::{
    CoreRegistryDataIntegrityShardingSubReport, CoreRegistryUpdateTxShardingSubReport,
};
use ex3_crypto::sha256;
use ex3_deposit_detector_public_types::{
    ConsumedDepositShardingSubReport, DepositDetectorDataIntegrityShardingSubReport,
};
use ex3_node_error::Error;
use ex3_node_types::balance_changed::BalanceChanged;
use ex3_node_types::block::{Block, BlockBody, BlockBodyExt, BlockHeader};
use ex3_node_types::number::Ex3Uint;
use ex3_node_types::transaction::{EncodedTransaction, Transaction, TransactionType};
use ex3_node_types::{
    AssetId, BalanceVaultSeqId, BlockHash, BlockHeight, CanisterId, PackageId, SecretVaultSeqId,
    Version, WalletRegistrySeqId,
};
use ex3_node_types::{MerkleNode, WalletRegisterId};
use ex3_secret_vault_public_types::{
    SecretUpdateTxShardingSubReport, SecretVaultDataIntegrityShardingSubReport,
};
use ex3_serde::bincode;
use ex3_wallet_registry_public_types::{
    WalletRegistryDataIntegrityShardingSubReport, WalletTxShardingSubReport,
};
use num_traits::{One, ToPrimitive};
use rs_merkle::algorithms::Sha256;
use rs_merkle::MerkleTree;

pub type Result<T> = result::Result<T, Error>;

pub struct SnapshotVaultIndex {
    wallet_registry_count: u32,
    secret_vault_count: u32,
    balance_vault_count: u32,
}

pub struct BlockBuilder {
    core_registry_canister_id: CanisterId,
    deposit_detector_canister_id: CanisterId,
    max_core_registry_update_tx_per_sharding_sub_report: usize,
    max_wallet_register_txs_per_sharding_sub_report: usize,
    max_secret_update_txs_per_sharding_sub_report: usize,
    max_balance_changed_records_per_sharding_sub_report: usize,
    max_withdrawal_txs_per_sharding_sub_report: usize,
    max_consumed_deposit_per_sharding_sub_report: usize,
    max_wallets_per_balance_vault_canister: usize,
    max_wallets_per_registry_canister: usize,
    max_wallets_per_secret_vault_canister: usize,
    start_wallet_id: usize,
}

mod balance_vault;
mod core_registry;
mod deposit_detector;
mod secret_vault;
mod wallet_registry;

impl BlockBuilder {
    pub fn new(
        core_registry_canister_id: CanisterId,
        deposit_detector_canister_id: CanisterId,
        max_core_registry_update_tx_per_sharding_sub_report: usize,
        max_wallet_register_txs_per_sharding_sub_report: usize,
        max_secret_update_txs_per_sharding_sub_report: usize,
        max_balance_changed_records_per_sharding_sub_report: usize,
        max_withdrawal_txs_per_sharding_sub_report: usize,
        max_consumed_deposit_per_sharding_sub_report: usize,
        max_wallets_per_balance_vault_canister: usize,
        max_wallets_per_registry_canister: usize,
        max_wallets_per_secret_vault_canister: usize,
        start_wallet_id: usize,
    ) -> Self {
        Self {
            core_registry_canister_id,
            deposit_detector_canister_id,
            max_core_registry_update_tx_per_sharding_sub_report,
            max_wallet_register_txs_per_sharding_sub_report,
            max_secret_update_txs_per_sharding_sub_report,
            max_balance_changed_records_per_sharding_sub_report,
            max_withdrawal_txs_per_sharding_sub_report,
            max_consumed_deposit_per_sharding_sub_report,
            max_wallets_per_balance_vault_canister,
            max_wallets_per_registry_canister,
            max_wallets_per_secret_vault_canister,
            start_wallet_id,
        }
    }

    pub fn build(
        &self,
        pre_block_hash: BlockHash,
        block_height: BlockHeight,
        last_process_p_id: PackageId,
        snapshot_cursor_pre_height: SnapshotCursor,
        txs: Vec<Transaction>,
        rejected_tx_indexes: Vec<u32>,
        balances_changed: Vec<(WalletRegisterId, Vec<(AssetId, BalanceChanged)>)>,
    ) -> Result<BlockWithShardingReports> {
        let (
            core_registry_tx_sharding_sub_reports,
            wallet_registry_sharding_sub_reports,
            secret_sharding_sub_reports,
            balance_sharding_sub_reports,
            deposit_detector_sharding_sub_reports,
            encoded_txs,
            tx_merkle_tree,
            rejected_tx_merkle_tree,
            balances_changed_merkle_tree,
            new_asset_count,
            new_market_count,
            new_wallet_count,
            total_accepted_withdrawal_count,
            total_rejected_withdrawal_count,
        ) = self.generate_sharding_data_sub_reports(
            &snapshot_cursor_pre_height,
            &txs,
            &rejected_tx_indexes,
            &balances_changed,
        );

        let (
            new_snapshot_vault_activated_report,
            id_cursor_report,
            core_registry_data_integrity_sharding_sub_report,
            wallet_registry_data_integrity_sharding_sub_reports,
            secret_vault_data_integrity_sharding_sub_reports,
            balance_vault_data_integrity_sharding_sub_reports,
            deposit_detector_data_integrity_sharding_sub_report,
            data_integrity_merkle_tree,
        ) = self.generate_data_integrity_sharding_sub_reports(
            &block_height,
            &core_registry_tx_sharding_sub_reports,
            &wallet_registry_sharding_sub_reports,
            &secret_sharding_sub_reports,
            &balance_sharding_sub_reports,
            &deposit_detector_sharding_sub_reports,
            &snapshot_cursor_pre_height,
            new_asset_count,
            new_market_count,
            new_wallet_count,
        );

        let block_body = BlockBody {
            txs: encoded_txs,
            data_integrity: data_integrity_merkle_tree.leaves().unwrap(),
            rejected_tx_indexes,
        };

        let block_header = BlockHeader {
            version: Version::V0.encode(),
            height: block_height,
            p_id: last_process_p_id,
            pre_block_hash,
            tx_root: tx_merkle_tree.root().map_or([0u8; 32], |root| root),
            state_changed_root: balances_changed_merkle_tree
                .root()
                .map_or([0u8; 32], |root| root),
            data_integrity_root: data_integrity_merkle_tree
                .root()
                .expect("data integrity merkle tree should be correct"),
            rejected_tx_root: rejected_tx_merkle_tree
                .root()
                .map_or([0u8; 32], |root| root),
        };

        let withdrawal_counter_report = WithdrawalCounterReport {
            accepted_withdrawal_count: total_accepted_withdrawal_count.try_into().unwrap(),
            rejected_withdrawal_count: total_rejected_withdrawal_count.try_into().unwrap(),
        };

        let block = Block {
            head: block_header.clone(),
            body: Some(block_body),
        };
        let block_bytes = bincode::serialize(&block).unwrap();
        let hash = cid::multihash::Code::Sha2_256.digest(&block_bytes);
        let block_cid = cid::Cid::new_v1(0x55, hash);

        let block_with_reports = BlockWithShardingReports {
            block: Arc::new(block),
            block_body_ext: Arc::new(BlockBodyExt { balances_changed }),
            consensus_report: CandidConsensusReport {
                head: block_header.into(),
                cid: block_cid.to_bytes(),
                activated_vault_report: new_snapshot_vault_activated_report,
                id_cursor_report,
                withdrawal_counter_report,
            },
            core_registry_reports: core_registry_tx_sharding_sub_reports,
            core_registry_data_integrity_report: core_registry_data_integrity_sharding_sub_report,
            wallet_registry_sharding_sub_reports,
            wallet_registry_data_integrity_reports:
                wallet_registry_data_integrity_sharding_sub_reports,
            secret_sharding_sub_reports,
            secret_vault_data_integrity_reports: secret_vault_data_integrity_sharding_sub_reports,
            balance_sharding_sub_reports,
            balance_vault_data_integrity_reports: balance_vault_data_integrity_sharding_sub_reports,
            deposit_detector_sharding_sub_reports,
            deposit_detector_data_integrity_report:
                deposit_detector_data_integrity_sharding_sub_report,
        };

        Ok(block_with_reports)
    }

    fn generate_sharding_data_sub_reports(
        &self,
        snapshot_vault_cursor_of_pre_height: &SnapshotCursor,
        txs: &Vec<Transaction>,
        rejected_tx_indexes: &Vec<u32>,
        balances_changed: &Vec<(WalletRegisterId, Vec<(AssetId, BalanceChanged)>)>,
    ) -> (
        Vec<CoreRegistryUpdateTxShardingSubReport>,
        BTreeMap<WalletRegistrySeqId, Vec<WalletTxShardingSubReport>>,
        BTreeMap<SecretVaultSeqId, Vec<SecretUpdateTxShardingSubReport>>,
        BTreeMap<
            BalanceVaultSeqId,
            (
                Vec<CandidBalanceShardingSubReport>,
                Vec<AcceptedWithdrawalShardingSubReport>,
                Vec<RejectedWithdrawalShardingSubReport>,
            ),
        >,
        Vec<ConsumedDepositShardingSubReport>,
        Vec<EncodedTransaction>,
        MerkleTree<Sha256>, // tx merkle tree
        MerkleTree<Sha256>, // rejected tx merkle tree
        MerkleTree<Sha256>,
        usize, // new asset count
        usize, // new market count
        usize, // new wallet count
        usize, // total accepted withdrawal count
        usize, // total rejected withdrawal count
    ) {
        let mut encoded_txs = vec![];
        let mut encoded_rejected_txs = vec![];
        let mut core_registry_update_txs = vec![];
        let mut wallet_tx_map = BTreeMap::new();
        let mut secret_vault_tx_map = BTreeMap::new();
        let mut balances_changed_map = BTreeMap::new();
        let mut balance_vault_accepted_withdrawal_tx_map = BTreeMap::new();
        let mut balance_vault_rejected_withdrawal_tx_map = BTreeMap::new();
        let mut deposit_txs = vec![];
        let mut new_asset_count = 0usize;
        let mut new_market_count = 0usize;
        let mut new_wallet_count = 0usize;
        let mut total_accepted_withdrawal_count = 0usize;
        let mut total_rejected_withdrawal_count = 0usize;
        let mut next_wallet_id = snapshot_vault_cursor_of_pre_height
            .max_wallet_id
            .clone()
            .map_or(self.start_wallet_id.clone().into(), |max_wallet_id| {
                WalletRegisterId::from(max_wallet_id + Ex3Uint::one())
            });
        let rejected_tx_indexes_map = rejected_tx_indexes
            .iter()
            .map(|index| (index.clone() as usize, ()))
            .collect::<BTreeMap<usize, ()>>();

        for (index, tx) in txs.iter().enumerate() {
            match tx.r#type {
                TransactionType::WalletRegister => {
                    let wallet_registry_seq_id =
                        self.get_wallet_registry_seq_id_by_wallet_id(&next_wallet_id);
                    let encoded_tx = EncodedTransaction::from(tx);
                    encoded_txs.push(encoded_tx.clone());
                    wallet_tx_map
                        .entry(wallet_registry_seq_id)
                        .or_insert(vec![])
                        .push((encoded_tx, index));
                    next_wallet_id += 1u8;
                    new_wallet_count += 1;
                }
                TransactionType::AssetAccountBinding | TransactionType::AssetAccountUnbinding => {
                    let wallet_registry_seq_id =
                        self.get_wallet_registry_seq_id_by_wallet_id(&tx.from);
                    let encoded_tx = EncodedTransaction::from(tx);
                    encoded_txs.push(encoded_tx.clone());
                    wallet_tx_map
                        .entry(wallet_registry_seq_id)
                        .or_insert(vec![])
                        .push((encoded_tx, index));
                }
                TransactionType::CreateApiSecret
                | TransactionType::DestroyApiSecret
                | TransactionType::ResetMainSecret => {
                    let secret_vault_seq_id = self.get_secret_vault_seq_id_by_wallet_id(&tx.from);
                    let encoded_tx = EncodedTransaction::from(tx);
                    encoded_txs.push(encoded_tx.clone());
                    secret_vault_tx_map
                        .entry(secret_vault_seq_id)
                        .or_insert(vec![])
                        .push((encoded_tx, index));
                }
                TransactionType::Withdrawal | TransactionType::ForceWithdrawal => {
                    let balance_vault_seq_id = self.get_balance_vault_seq_id_by_wallet_id(&tx.from);
                    let encoded_tx = EncodedTransaction::from(tx);

                    if rejected_tx_indexes_map.contains_key(&index) {
                        encoded_rejected_txs.push(encoded_tx.clone());
                        balance_vault_rejected_withdrawal_tx_map
                            .entry(balance_vault_seq_id)
                            .or_insert(vec![])
                            .push((encoded_tx, index));
                        total_rejected_withdrawal_count += 1;
                    } else {
                        encoded_txs.push(encoded_tx.clone());
                        balance_vault_accepted_withdrawal_tx_map
                            .entry(balance_vault_seq_id)
                            .or_insert(vec![])
                            .push((encoded_tx, index));
                        total_accepted_withdrawal_count += 1;
                    }
                }
                TransactionType::RegisterAsset
                | TransactionType::UpdateAssetWithdrawalFeeTo
                | TransactionType::UpdateChainConfirmationTimes
                | TransactionType::UpdateGlobalWithdrawalFeeTo
                | TransactionType::RegisterSpotMarket
                | TransactionType::UpdateSpotMarketTradingSettings
                | TransactionType::UpdateSpotMarketInitialFeeTo
                | TransactionType::UpdateSpotMarketInitialTradingFee
                | TransactionType::UpdateSpotMarketFeeTo
                | TransactionType::UpdateSpotMarketTradingFee => {
                    if tx.r#type == TransactionType::RegisterAsset {
                        new_asset_count += 1;
                    } else if tx.r#type == TransactionType::RegisterSpotMarket {
                        new_market_count += 1;
                    }
                    let encoded_tx = EncodedTransaction::from(tx);
                    encoded_txs.push(encoded_tx.clone());
                    core_registry_update_txs.push((encoded_tx, index));
                }
                TransactionType::Deposit => {
                    let encoded_tx = EncodedTransaction::from(tx);
                    encoded_txs.push(encoded_tx.clone());
                    deposit_txs.push((encoded_tx, index));
                }
                _ => {
                    let encoded_tx = EncodedTransaction::from(tx);
                    encoded_txs.push(encoded_tx.clone());
                }
            }
        }

        let tx_merkle_tree = self.build_tx_merkle_tree(&encoded_txs);
        let rejected_txs_merkle_tree = self.build_rejected_txs_merkle_tree(&encoded_rejected_txs);
        let balances_changed_merkle_tree =
            self.build_balances_changed_merkle_tree(&balances_changed);

        for (index, (wallet_id, asset_changes)) in balances_changed.iter().enumerate() {
            let balance_vault_seq_id = self.get_balance_vault_seq_id_by_wallet_id(&wallet_id);
            balances_changed_map
                .entry(balance_vault_seq_id)
                .or_insert(vec![])
                .push(((wallet_id.clone(), asset_changes.clone()), index));
        }

        let core_registry_tx_sharding_sub_reports = self
            .generate_core_registry_update_tx_sharding_sub_reports(
                core_registry_update_txs,
                &tx_merkle_tree,
            );

        let wallet_register_tx_sharding_sub_reports = self.generate_wallet_tx_sharding_sub_reports(
            wallet_tx_map,
            &tx_merkle_tree,
            snapshot_vault_cursor_of_pre_height,
        );

        let secret_update_tx_sharding_sub_reports = self
            .generate_secret_update_tx_sharding_sub_reports(
                secret_vault_tx_map,
                &tx_merkle_tree,
                snapshot_vault_cursor_of_pre_height,
            );

        let consumed_deposit_sharding_sub_reports =
            self.generate_consumed_deposit_sharding_sub_reports(deposit_txs, &tx_merkle_tree);

        let balance_sharding_sub_reports = self.generate_balance_sharding_sub_reports(
            balances_changed_map,
            &balances_changed_merkle_tree,
            balance_vault_accepted_withdrawal_tx_map,
            &tx_merkle_tree,
            balance_vault_rejected_withdrawal_tx_map,
            &rejected_txs_merkle_tree,
            snapshot_vault_cursor_of_pre_height,
        );

        (
            core_registry_tx_sharding_sub_reports,
            wallet_register_tx_sharding_sub_reports,
            secret_update_tx_sharding_sub_reports,
            balance_sharding_sub_reports,
            consumed_deposit_sharding_sub_reports,
            encoded_txs,
            tx_merkle_tree,
            rejected_txs_merkle_tree,
            balances_changed_merkle_tree,
            new_asset_count,
            new_market_count,
            new_wallet_count,
            total_accepted_withdrawal_count,
            total_rejected_withdrawal_count,
        )
    }

    /// Generate data integrity sharding data integrity sub reports
    fn generate_data_integrity_sharding_sub_reports(
        &self,
        current_block_height: &BlockHeight,
        core_registry_tx_sharding_sub_reports: &Vec<CoreRegistryUpdateTxShardingSubReport>,
        wallet_register_tx_sharding_sub_reports: &BTreeMap<
            WalletRegistrySeqId,
            Vec<WalletTxShardingSubReport>,
        >,
        secret_update_tx_sharding_sub_reports: &BTreeMap<
            SecretVaultSeqId,
            Vec<SecretUpdateTxShardingSubReport>,
        >,
        balance_sharding_sub_reports: &BTreeMap<
            BalanceVaultSeqId,
            (
                Vec<CandidBalanceShardingSubReport>,
                Vec<AcceptedWithdrawalShardingSubReport>,
                Vec<RejectedWithdrawalShardingSubReport>,
            ),
        >,
        deposit_detector_sharding_sub_reports: &Vec<ConsumedDepositShardingSubReport>,
        snapshot_vault_cursor_of_pre_height: &SnapshotCursor,
        new_asset_count: usize,
        new_market_count: usize,
        new_wallet_count: usize,
    ) -> (
        NewSnapshotVaultActivatedReport,
        IdCursorReport,
        CoreRegistryDataIntegrityShardingSubReport,
        BTreeMap<WalletRegistrySeqId, WalletRegistryDataIntegrityShardingSubReport>,
        BTreeMap<SecretVaultSeqId, SecretVaultDataIntegrityShardingSubReport>,
        BTreeMap<BalanceVaultSeqId, BalanceDataIntegrityShardingSubReport>,
        DepositDetectorDataIntegrityShardingSubReport,
        MerkleTree<Sha256>,
    ) {
        let snapshot_vault_index = SnapshotVaultIndex::new(
            wallet_register_tx_sharding_sub_reports
                .len()
                .try_into()
                .unwrap(),
            secret_update_tx_sharding_sub_reports
                .len()
                .try_into()
                .unwrap(),
            balance_sharding_sub_reports.len().try_into().unwrap(),
        );

        let core_registry_data_integrity = self.generate_core_registry_data_integrity(
            current_block_height,
            core_registry_tx_sharding_sub_reports,
            &snapshot_vault_index,
        );

        let (wallet_registry_data_integrity, new_activated_wallet_registries) = self
            .generate_wallet_registry_data_integrity(
                current_block_height,
                &wallet_register_tx_sharding_sub_reports,
                snapshot_vault_cursor_of_pre_height,
                &snapshot_vault_index,
            );

        let (secret_vault_data_integrity, new_activated_secret_vaults) = self
            .generate_secret_vault_data_integrity(
                current_block_height,
                &secret_update_tx_sharding_sub_reports,
                snapshot_vault_cursor_of_pre_height,
                &snapshot_vault_index,
            );

        let (balance_vault_data_integrity, new_activated_balance_vaults) = self
            .generate_balance_vault_data_integrity(
                current_block_height,
                &balance_sharding_sub_reports,
                snapshot_vault_cursor_of_pre_height,
                &snapshot_vault_index,
            );

        let deposit_detector_data_integrity = self.generate_deposit_detector_data_integrity(
            current_block_height,
            deposit_detector_sharding_sub_reports,
            &snapshot_vault_index,
        );

        let id_cursor_report = IdCursorReport {
            wallet_id_cursor_increase: new_wallet_count.try_into().unwrap(),
            asset_id_cursor_increase: new_asset_count.try_into().unwrap(),
            market_id_cursor_increase: new_market_count.try_into().unwrap(),
        };

        let mut data_integrity = vec![];

        data_integrity.push(core_registry_data_integrity.clone());
        data_integrity.extend(wallet_registry_data_integrity.values());
        data_integrity.extend(secret_vault_data_integrity.values());
        data_integrity.extend(balance_vault_data_integrity.values());
        data_integrity.push(deposit_detector_data_integrity.clone());
        data_integrity.push(sha256(
            id_cursor_report.wallet_id_cursor_increase.to_le_bytes(),
        ));
        data_integrity.push(sha256(
            id_cursor_report.asset_id_cursor_increase.to_le_bytes(),
        ));
        data_integrity.push(sha256(
            id_cursor_report.market_id_cursor_increase.to_le_bytes(),
        ));

        let data_integrity_merkle_tree = self.build_data_integrity_merkle_tree(&data_integrity);

        let (
            core_registry_data_integrity_report,
            wallet_registry_data_integrity_reports,
            secret_vault_data_integrity_reports,
            balance_vault_data_integrity_reports,
            deposit_detector_data_integrity_report,
        ) = self.generate_data_integrity_reports(
            &snapshot_vault_index,
            &core_registry_data_integrity,
            &core_registry_tx_sharding_sub_reports,
            &wallet_registry_data_integrity,
            &wallet_register_tx_sharding_sub_reports,
            &secret_vault_data_integrity,
            &secret_update_tx_sharding_sub_reports,
            &balance_vault_data_integrity,
            &balance_sharding_sub_reports,
            &deposit_detector_data_integrity,
            &deposit_detector_sharding_sub_reports,
            &data_integrity_merkle_tree,
        );

        let new_snapshot_vault_activated_report = NewSnapshotVaultActivatedReport {
            wallet_registries: new_activated_wallet_registries,
            secret_vaults: new_activated_secret_vaults,
            balance_vaults: new_activated_balance_vaults,
        };

        (
            new_snapshot_vault_activated_report,
            id_cursor_report,
            core_registry_data_integrity_report,
            wallet_registry_data_integrity_reports,
            secret_vault_data_integrity_reports,
            balance_vault_data_integrity_reports,
            deposit_detector_data_integrity_report,
            data_integrity_merkle_tree,
        )
    }

    /// Generate data integrity reports
    fn generate_data_integrity_reports(
        &self,
        snapshot_vault_index: &SnapshotVaultIndex,
        core_registry_data_integrity: &MerkleNode,
        core_registry_sharding_sub_reports: &Vec<CoreRegistryUpdateTxShardingSubReport>,
        wallet_registry_data_integrity: &BTreeMap<WalletRegistrySeqId, MerkleNode>,
        wallet_registry_sharding_sub_reports: &BTreeMap<
            WalletRegistrySeqId,
            Vec<WalletTxShardingSubReport>,
        >,
        secret_vault_data_integrity: &BTreeMap<SecretVaultSeqId, MerkleNode>,
        secret_vault_sharding_sub_reports: &BTreeMap<
            SecretVaultSeqId,
            Vec<SecretUpdateTxShardingSubReport>,
        >,
        balance_vault_data_integrity: &BTreeMap<BalanceVaultSeqId, MerkleNode>,
        balance_sharding_sub_reports: &BTreeMap<
            BalanceVaultSeqId,
            (
                Vec<CandidBalanceShardingSubReport>,
                Vec<AcceptedWithdrawalShardingSubReport>,
                Vec<RejectedWithdrawalShardingSubReport>,
            ),
        >,
        deposit_detector_data_integrity: &MerkleNode,
        deposit_detector_sharding_sub_reports: &Vec<ConsumedDepositShardingSubReport>,
        data_integrity_merkle_tree: &MerkleTree<Sha256>,
    ) -> (
        CoreRegistryDataIntegrityShardingSubReport,
        BTreeMap<WalletRegistrySeqId, WalletRegistryDataIntegrityShardingSubReport>,
        BTreeMap<SecretVaultSeqId, SecretVaultDataIntegrityShardingSubReport>,
        BTreeMap<BalanceVaultSeqId, BalanceDataIntegrityShardingSubReport>,
        DepositDetectorDataIntegrityShardingSubReport,
    ) {
        let total_merkle_node_count = data_integrity_merkle_tree.leaves_len();

        /* core registry report generate */
        let core_registry_index = snapshot_vault_index.get_core_registry_index();
        let core_registry_merkle_proof =
            data_integrity_merkle_tree.proof(&[core_registry_index as usize]);
        let total_core_registry_update_tx_count = core_registry_sharding_sub_reports
            .iter()
            .fold(0, |acc, report| acc + report.txs.len());

        let core_registry_report = CoreRegistryDataIntegrityShardingSubReport {
            merkle_proof: core_registry_merkle_proof.to_bytes(),
            data_digest: core_registry_data_integrity.clone(),
            merkle_node_count: total_merkle_node_count.try_into().unwrap(),
            index: core_registry_index,
            total_update_tx_count: total_core_registry_update_tx_count.try_into().unwrap(),
        };

        /* wallet registry report generate */
        let wallet_registry_reports = wallet_registry_data_integrity
            .iter()
            .map(|(seq_id, data_integrity)| {
                let vault_index = snapshot_vault_index.get_wallet_registry_index(seq_id);
                let sub_reports = wallet_registry_sharding_sub_reports.get(seq_id);
                let merkle_proof = data_integrity_merkle_tree.proof(&[vault_index as usize]);
                let total_register_tx_count = sub_reports.map_or(0, |reports| {
                    reports.iter().map(|report| report.txs.len()).sum::<usize>()
                });
                let report = WalletRegistryDataIntegrityShardingSubReport {
                    merkle_proof: merkle_proof.to_bytes(),
                    data_digest: data_integrity.clone(),
                    index: vault_index,
                    merkle_node_count: total_merkle_node_count.try_into().unwrap(),
                    total_register_tx_count: total_register_tx_count.try_into().unwrap(),
                };
                (seq_id.clone(), report)
            })
            .collect();

        let secret_vault_reports = secret_vault_data_integrity
            .iter()
            .map(|(seq_id, data_integrity)| {
                let vault_index = snapshot_vault_index.get_secret_vault_index(seq_id);
                let merkle_proof = data_integrity_merkle_tree.proof(&[vault_index as usize]);
                let total_update_tx_count = secret_vault_sharding_sub_reports
                    .get(seq_id)
                    .map_or(0, |reports| {
                        reports.iter().map(|report| report.txs.len()).sum::<usize>()
                    });
                let report = SecretVaultDataIntegrityShardingSubReport {
                    merkle_proof: merkle_proof.to_bytes(),
                    data_digest: data_integrity.clone(),
                    merkle_node_count: total_merkle_node_count.try_into().unwrap(),
                    index: vault_index,
                    total_update_tx_count: total_update_tx_count.try_into().unwrap(),
                };
                (seq_id.clone(), report)
            })
            .collect();

        let balance_vault_reports = balance_vault_data_integrity
            .iter()
            .map(|(seq_id, data_integrity)| {
                let vault_index = snapshot_vault_index.get_balance_vault_index(seq_id);
                let merkle_proof = data_integrity_merkle_tree.proof(&[vault_index as usize]);

                let (
                    wallet_count_balance_changed,
                    accepted_withdrawal_count,
                    rejected_withdrawal_count,
                ) = balance_sharding_sub_reports
                    .get(seq_id)
                    .map_or((0, 0, 0), |reports| {
                        let wallet_count_balance_changed = reports
                            .0
                            .iter()
                            .map(|report| report.balances_changed.0.len())
                            .sum();
                        let accepted_withdrawal_count = reports
                            .1
                            .iter()
                            .map(|report| report.withdrawals.len())
                            .sum();
                        let rejected_withdrawal_count = reports
                            .2
                            .iter()
                            .map(|report| report.withdrawals.len())
                            .sum();
                        (
                            wallet_count_balance_changed,
                            accepted_withdrawal_count,
                            rejected_withdrawal_count,
                        )
                    });

                let report = BalanceDataIntegrityShardingSubReport {
                    merkle_proof: merkle_proof.to_bytes(),
                    data_digest: data_integrity.clone(),
                    wallet_count_balance_changed: wallet_count_balance_changed.try_into().unwrap(),
                    accepted_withdrawal_count: accepted_withdrawal_count.try_into().unwrap(),
                    rejected_withdrawal_count: rejected_withdrawal_count.try_into().unwrap(),
                    index: vault_index.try_into().unwrap(),
                    merkle_node_count: total_merkle_node_count.try_into().unwrap(),
                };
                (seq_id.clone(), report)
            })
            .collect();

        let deposit_detector_index = snapshot_vault_index.get_deposit_detector_index();
        let deposit_detector_merkle_proof =
            data_integrity_merkle_tree.proof(&[deposit_detector_index as usize]);
        let total_consumed_deposit_count = deposit_detector_sharding_sub_reports
            .iter()
            .fold(0, |acc, report| acc + report.deposit_txs.len());
        let deposit_detector_report = DepositDetectorDataIntegrityShardingSubReport {
            merkle_proof: deposit_detector_merkle_proof.to_bytes(),
            data_digest: deposit_detector_data_integrity.clone(),
            merkle_node_count: total_merkle_node_count.try_into().unwrap(),
            index: deposit_detector_index,
            total_deposit_tx_count: total_consumed_deposit_count.try_into().unwrap(),
        };

        (
            core_registry_report,
            wallet_registry_reports,
            secret_vault_reports,
            balance_vault_reports,
            deposit_detector_report,
        )
    }

    fn build_tx_merkle_tree(&self, txs: &Vec<EncodedTransaction>) -> MerkleTree<Sha256> {
        let nodes = txs
            .iter()
            .map(|tx| sha256(tx.as_ref()))
            .collect::<Vec<MerkleNode>>();
        MerkleTree::<Sha256>::from_leaves(&nodes)
    }

    fn build_rejected_txs_merkle_tree(
        &self,
        rejected_txs: &Vec<EncodedTransaction>,
    ) -> MerkleTree<Sha256> {
        let nodes = rejected_txs
            .iter()
            .map(|tx| sha256(tx.as_ref()))
            .collect::<Vec<MerkleNode>>();
        MerkleTree::<Sha256>::from_leaves(&nodes)
    }

    pub fn build_data_integrity_merkle_tree(
        &self,
        data_integrity: &Vec<MerkleNode>,
    ) -> MerkleTree<Sha256> {
        MerkleTree::<Sha256>::from_leaves(data_integrity)
    }
}

impl SnapshotVaultIndex {
    pub fn new(
        wallet_registry_count: u32,
        secret_vault_count: u32,
        balance_vault_count: u32,
    ) -> Self {
        Self {
            wallet_registry_count,
            secret_vault_count,
            balance_vault_count,
        }
    }

    pub fn get_total_vaults_count(&self) -> u32 {
        self.wallet_registry_count.clone()
            + self.secret_vault_count.clone()
            + self.balance_vault_count.clone()
            + 2u32
    }

    pub fn get_total_merkle_node_count(&self) -> u32 {
        // total count of the canisters, include
        //   core registry, wallet registry, balance vault, secret vault, deposit detector
        //   + 3 (last merkle node is the new wallet count,new asset count, new market count)
        self.get_total_vaults_count() + 3u32
    }

    pub fn get_core_registry_index(&self) -> u32 {
        0u32
    }

    pub fn get_wallet_registry_index(&self, wallet_registry_seq_id: &WalletRegistrySeqId) -> u32 {
        wallet_registry_seq_id.to_u32().unwrap() + 1u32
    }

    pub fn get_secret_vault_index(&self, secret_vault_seq_id: &SecretVaultSeqId) -> u32 {
        secret_vault_seq_id.to_u32().unwrap() + self.wallet_registry_count.clone() + 1u32
    }

    pub fn get_balance_vault_index(&self, balance_vault_seq_id: &BalanceVaultSeqId) -> u32 {
        balance_vault_seq_id.to_u32().unwrap()
            + self.wallet_registry_count.clone()
            + self.secret_vault_count.clone()
            + 1u32
    }

    pub fn get_deposit_detector_index(&self) -> u32 {
        self.wallet_registry_count.clone()
            + self.secret_vault_count.clone()
            + self.balance_vault_count.clone()
            + 1u32
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_generate_merkle_tree_with_empty_nodes() {
        let block_builder = BlockBuilder::new(
            CanisterId::from_text("a4tbr-q4aaa-aaaaa-qaafq-cai").unwrap(),
            CanisterId::from_text("a4tbr-q4aaa-aaaaa-qaafq-cai").unwrap(),
            100,
            100,
            100,
            100,
            100,
            100,
            100,
            100,
            100,
            100,
        );
        let merkle_tree = block_builder.build_rejected_txs_merkle_tree(&vec![]);
        assert_eq!(merkle_tree.leaves_len(), 0);
        assert_eq!(merkle_tree.root().is_none(), true);
    }
}