use std::ops::Add;

use candid::{CandidType, Deserialize, Nat};
use ex3_canister_constant::{START_ASSET_ID, START_MARKET_ID, START_WALLET_ID};
use ex3_crypto::sha256;
use ex3_node_types::block::{BlockHead, CandidBlockHead};
use ex3_node_types::{
    AssetId, CandidAssetId, CandidMarketId, CandidWalletRegisterId, CanisterId, MarketId,
    WalletRegisterId,
};
use num_bigint::BigUint;
use num_traits::{CheckedSub, One, ToPrimitive, Zero};
use rs_merkle::algorithms::Sha256;
use rs_merkle::MerkleProof;
use serde::Serialize;
use serde_bytes::ByteBuf;

pub mod error;

type BlockchainArchiveCanister = CanisterId;

#[derive(CandidType, Deserialize, Clone, Debug, PartialEq, Eq)]
pub enum BlockHeadResponse {
    Ok(Option<CandidBlockHead>),
    Forward(BlockchainArchiveCanister),
}

#[derive(CandidType, Deserialize, Debug, Clone)]
pub struct CandidConsensusReport {
    pub head: CandidBlockHead,
    pub merkle_proof: ByteBuf,
    pub activated_vault_report: NewSnapshotVaultActivatedReport,
    pub id_cursor_report: IdCursorReport,
    pub withdrawal_counter_report: WithdrawalCounterReport,
}

#[derive(Deserialize, Serialize, Debug, Clone, PartialEq, Eq)]
pub struct ConsensusReport {
    pub head: BlockHead,
    pub merkle_proof: ByteBuf,
    pub activated_vault_report: NewSnapshotVaultActivatedReport,
    pub id_cursor_report: IdCursorReport,
    pub withdrawal_counter_report: WithdrawalCounterReport,
}

impl ConsensusReport {
    pub fn get_new_vault_increase_count(&self) -> usize {
        self.activated_vault_report.wallet_registries.len()
            + self.activated_vault_report.secret_vaults.len()
            + self.activated_vault_report.balance_vaults.len()
    }
    pub fn merkle_verify(&self, total_vault_count: BigUint) -> bool {
        let merkle_proof: MerkleProof<Sha256> =
            MerkleProof::from_bytes(self.merkle_proof.as_ref()).unwrap();

        // The increase wallet count is the last node in the merkle tree
        let mut indexes = vec![];
        let mut leaves = vec![];

        self.activated_vault_report
            .wallet_registries
            .iter()
            .for_each(|w| {
                indexes.push(w.canister_index);
                leaves.push(w.hash_of_sharding);
            });

        self.activated_vault_report
            .secret_vaults
            .iter()
            .for_each(|sv| {
                indexes.push(sv.canister_index);
                leaves.push(sv.hash_of_sharding);
            });

        self.activated_vault_report
            .balance_vaults
            .iter()
            .for_each(|bv| {
                indexes.push(bv.canister_index);
                leaves.push(bv.hash_of_sharding);
            });

        let total_vault_count = total_vault_count
            .to_usize()
            .expect("total_vault_count overflow");

        let bytes_of_increase_wallet_count = self
            .id_cursor_report
            .wallet_id_cursor_increase
            .to_be_bytes();
        indexes.push(total_vault_count);
        leaves.push(sha256!(bytes_of_increase_wallet_count));
        merkle_proof.verify(
            self.head.data_integrity_root,
            &indexes,
            &leaves,
            total_vault_count + 3,
        )
    }
}

#[derive(CandidType, Deserialize, Serialize, Default, Debug, Clone, PartialEq, Eq)]
pub struct NewSnapshotVaultActivatedReport {
    pub wallet_registries: Vec<ActivatedWalletRegistry>,
    pub secret_vaults: Vec<ActivatedSecretVault>,
    pub balance_vaults: Vec<ActivatedBalanceVault>,
}

#[derive(CandidType, Deserialize, Serialize, Default, Debug, Clone, PartialEq, Eq)]
pub struct IdCursorReport {
    pub wallet_id_cursor_increase: usize,
    pub asset_id_cursor_increase: usize,
    pub market_id_cursor_increase: usize,
}

#[derive(CandidType, Deserialize, Serialize, Default, Debug, Clone, PartialEq, Eq)]
pub struct WithdrawalCounterReport {
    pub accepted_withdrawal_count: usize,
    pub rejected_withdrawal_count: usize,
}

pub type ActivatedWalletRegistry = ActivatedSnapshotVault;
pub type ActivatedSecretVault = ActivatedSnapshotVault;
pub type ActivatedBalanceVault = ActivatedSnapshotVault;

#[derive(CandidType, Deserialize, Serialize, Debug, Clone, PartialEq, Eq)]
pub struct ActivatedSnapshotVault {
    pub canister_seq_id: u64,
    pub canister_total_count: usize,
    pub canister_index: usize,
    pub hash_of_sharding: [u8; 32],
}

#[derive(CandidType, Deserialize, Clone, Debug, PartialEq, Eq)]
pub struct CandidSnapshotCursor {
    /// The snapshot max activated wallet registry seq id
    pub max_activated_wallet_registry_seq_id: u64,

    /// The snapshot max activated secret vault seq id
    pub max_activated_secret_vault_seq_id: u64,

    /// The snapshot max activated balance vault seq id
    pub max_activated_balance_vault_seq_id: u64,

    /// The snapshot vault count
    /// Includes core registry,wallet registry,secret vault, and balance vault.
    pub vault_count: Nat,

    /// The snapshot max wallet id
    pub max_wallet_id: Option<CandidWalletRegisterId>,

    /// The snapshot max asset id
    pub max_asset_id: Option<CandidAssetId>,

    /// The snapshot max market id
    pub max_market_id: Option<CandidMarketId>,
}

#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)]
pub struct SnapshotCursor {
    /// The snapshot max activated wallet registry seq id
    pub max_activated_wallet_registry_seq_id: u64,

    /// The snapshot max activated secret vault seq id
    pub max_activated_secret_vault_seq_id: u64,

    /// The snapshot max activated balance vault seq id
    pub max_activated_balance_vault_seq_id: u64,

    /// The snapshot vault count
    /// Includes core registry,wallet registry,secret vault, and balance vault.
    pub vault_count: BigUint,

    /// The snapshot max wallet id
    pub max_wallet_id: Option<WalletRegisterId>,

    /// The snapshot max asset id
    pub max_asset_id: Option<AssetId>,

    /// The snapshot max market id
    pub max_market_id: Option<MarketId>,
}

impl SnapshotCursor {
    fn apply_activated_wallet_registry(
        &mut self,
        activated_wallet_registry: &ActivatedWalletRegistry,
    ) {
        if activated_wallet_registry.canister_seq_id > self.max_activated_wallet_registry_seq_id {
            self.max_activated_wallet_registry_seq_id = activated_wallet_registry.canister_seq_id;
        }
    }

    fn apply_activated_secret_vault(&mut self, activated_secret_vault: &ActivatedSecretVault) {
        if activated_secret_vault.canister_seq_id > self.max_activated_secret_vault_seq_id {
            self.max_activated_secret_vault_seq_id = activated_secret_vault.canister_seq_id;
        }
    }

    fn apply_activated_balance_vault(&mut self, activated_balance_vault: &ActivatedBalanceVault) {
        if activated_balance_vault.canister_seq_id > self.max_activated_balance_vault_seq_id {
            self.max_activated_balance_vault_seq_id = activated_balance_vault.canister_seq_id;
        }
    }

    fn apply_activated_vault_report(
        &mut self,
        activated_vault_report: &NewSnapshotVaultActivatedReport,
    ) {
        for activated_wallet_registry in &activated_vault_report.wallet_registries {
            self.apply_activated_wallet_registry(activated_wallet_registry);
        }
        for activated_secret_vault in &activated_vault_report.secret_vaults {
            self.apply_activated_secret_vault(activated_secret_vault);
        }
        for activated_balance_vault in &activated_vault_report.balance_vaults {
            self.apply_activated_balance_vault(activated_balance_vault);
        }
        self.vault_count += BigUint::from(activated_vault_report.wallet_registries.len());
        self.vault_count += BigUint::from(activated_vault_report.secret_vaults.len());
        self.vault_count += BigUint::from(activated_vault_report.balance_vaults.len());
    }

    fn apply_id_cursor_report(&mut self, id_cursor_report: &IdCursorReport) {
        let wallet_id_cursor_increase = id_cursor_report.wallet_id_cursor_increase.clone();
        let current_max_wallet_id = self.max_wallet_id.clone();
        let new_max_wallet_id = current_max_wallet_id.map_or(
            wallet_id_cursor_increase
                .clone()
                .add(&BigUint::from(START_WALLET_ID))
                .checked_sub(&BigUint::one())
                .expect("Invalid max wallet id"),
            |id| id + wallet_id_cursor_increase,
        );
        self.max_wallet_id = Some(new_max_wallet_id);

        let asset_id_cursor_increase = id_cursor_report.asset_id_cursor_increase.clone();
        let current_max_asset_id = self.max_asset_id.clone();
        let new_max_asset_id = current_max_asset_id.map_or(
            asset_id_cursor_increase
                .clone()
                .add(&BigUint::from(START_ASSET_ID))
                .checked_sub(&BigUint::one())
                .expect("Invalid max asset id"),
            |id| id + asset_id_cursor_increase,
        );

        self.max_asset_id = Some(new_max_asset_id);

        let market_id_cursor_increase = id_cursor_report.market_id_cursor_increase.clone();
        let current_max_market_id = self.max_market_id.clone();
        let new_max_market_id = current_max_market_id.map_or(
            market_id_cursor_increase
                .clone()
                .add(&BigUint::from(START_MARKET_ID))
                .checked_sub(&BigUint::one())
                .expect("Invalid max market id"),
            |id| id + market_id_cursor_increase,
        );

        self.max_market_id = Some(new_max_market_id);
    }

    pub fn apply_consensus_report(&mut self, consensus_report: &ConsensusReport) {
        self.apply_activated_vault_report(&consensus_report.activated_vault_report);
        self.apply_id_cursor_report(&consensus_report.id_cursor_report);
    }
}

impl Default for SnapshotCursor {
    fn default() -> Self {
        SnapshotCursor {
            max_activated_wallet_registry_seq_id: 0,
            max_activated_secret_vault_seq_id: 0,
            max_activated_balance_vault_seq_id: 0,
            vault_count: BigUint::zero(),
            max_wallet_id: None,
            max_asset_id: None,
            max_market_id: None,
        }
    }
}

impl From<CandidSnapshotCursor> for SnapshotCursor {
    fn from(candid_snapshot_cursor: CandidSnapshotCursor) -> Self {
        Self {
            max_activated_wallet_registry_seq_id: candid_snapshot_cursor
                .max_activated_wallet_registry_seq_id,
            max_activated_secret_vault_seq_id: candid_snapshot_cursor
                .max_activated_secret_vault_seq_id,
            max_activated_balance_vault_seq_id: candid_snapshot_cursor
                .max_activated_balance_vault_seq_id,
            vault_count: candid_snapshot_cursor.vault_count.into(),
            max_wallet_id: candid_snapshot_cursor
                .max_wallet_id
                .map(|wallet_id| wallet_id.0),
            max_asset_id: candid_snapshot_cursor
                .max_asset_id
                .map(|asset_id| asset_id.0),
            max_market_id: candid_snapshot_cursor
                .max_market_id
                .map(|market_id| market_id.0),
        }
    }
}

impl From<SnapshotCursor> for CandidSnapshotCursor {
    fn from(snapshot_cursor: SnapshotCursor) -> Self {
        Self {
            max_activated_wallet_registry_seq_id: snapshot_cursor
                .max_activated_wallet_registry_seq_id,
            max_activated_secret_vault_seq_id: snapshot_cursor.max_activated_secret_vault_seq_id,
            max_activated_balance_vault_seq_id: snapshot_cursor.max_activated_balance_vault_seq_id,
            vault_count: snapshot_cursor.vault_count.into(),
            max_wallet_id: snapshot_cursor.max_wallet_id.map(Into::into),
            max_asset_id: snapshot_cursor.max_asset_id.map(Into::into),
            max_market_id: snapshot_cursor.max_market_id.map(Into::into),
        }
    }
}

impl From<CandidConsensusReport> for ConsensusReport {
    fn from(candid_consensus_report: CandidConsensusReport) -> Self {
        Self {
            head: candid_consensus_report.head.into(),
            merkle_proof: candid_consensus_report.merkle_proof.into(),
            activated_vault_report: candid_consensus_report.activated_vault_report,
            id_cursor_report: candid_consensus_report.id_cursor_report,
            withdrawal_counter_report: candid_consensus_report.withdrawal_counter_report,
        }
    }
}
