use std::collections::BTreeMap;
use std::ops::{Add, Sub};
use std::sync::Arc;

use ex3_balance_vault_public_types::{
    BalanceDataIntegrityShardingSubReport, CandidBalanceShardingSubReport,
    CandidWithdrawalShardingSubReport,
};
use ex3_blockchain_public_types::{BlockHead, CandidConsensusReport};
use ex3_canister_types::report::DataIntegrityShardingSubReport;
use ex3_core_registry_public_types::CandidCoreRegistryUpdateTxShardingSubReport;
use ex3_secret_vault_public_types::CandidSecretUpdateTxShardingSubReport;
use ex3_serde::bincode::{deserialize, serialize};
use ex3_wallet_registry_public_types::CandidWalletRegisterTxShardingSubReport;
use serde::{Deserialize, Serialize};
use serde_bytes::ByteBuf;

use ex3_node_error::OtherError;

use crate::transaction::EncodedTransaction;
use crate::{AssetAmount, AssetId, BlockHeight, MerkleNode, VaultSeqId, WalletRegisterId};

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct BlockHeadExt {
    pub height: BlockHeight,
    pub validate_on_chain: bool,
    pub sharding_reported: bool,
    pub has_local_body: bool,
}

impl BlockHeadExt {
    pub fn encode(&self) -> Vec<u8> {
        serialize(&(
            &self.validate_on_chain,
            &self.sharding_reported,
            &self.has_local_body,
            &self.height,
        ))
        .unwrap()
    }

    pub fn decode(bytes: &[u8]) -> Result<Self, OtherError> {
        let (validate_on_chain, has_been_reported, has_local_body, height) =
            deserialize(bytes).map_err(|e| OtherError::new(format!("{:?}", e)))?;
        Ok(Self {
            validate_on_chain,
            sharding_reported: has_been_reported,
            has_local_body,
            height,
        })
    }
}

impl TryFrom<&[u8]> for BlockHeadExt {
    type Error = OtherError;

    fn try_from(bytes: &[u8]) -> Result<Self, Self::Error> {
        Self::decode(bytes)
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Block {
    pub head: BlockHead,
    pub body: Option<BlockBody>,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct BlockBody {
    pub txs: Vec<EncodedTransaction>,
    pub balances_changed: Vec<(WalletRegisterId, Vec<(AssetId, BalanceChanged)>)>,
    pub data_integrity: Vec<MerkleNode>,
}

impl BlockBody {
    pub fn encode(&self) -> Vec<u8> {
        serialize(&(&self.txs, &self.balances_changed)).unwrap()
    }
}

impl TryFrom<&[u8]> for BlockBody {
    type Error = OtherError;

    fn try_from(bytes: &[u8]) -> Result<Self, Self::Error> {
        let (txs, balances_changed, data_integrity) =
            deserialize(bytes).map_err(|e| OtherError::new(format!("{:?}", e)))?;
        Ok(Self {
            txs,
            balances_changed,
            data_integrity,
        })
    }
}

#[derive(Deserialize, Serialize, Debug, Clone, PartialEq, Eq)]
pub struct EncodedBlockBody(pub ByteBuf);

impl From<BlockBody> for EncodedBlockBody {
    fn from(body: BlockBody) -> Self {
        let bytes = serialize(&body.encode()).unwrap();
        EncodedBlockBody(ByteBuf::from(bytes))
    }
}

impl From<EncodedBlockBody> for BlockBody {
    fn from(encoded_body: EncodedBlockBody) -> Self {
        encoded_body.0.as_slice().try_into().unwrap()
    }
}

#[derive(Deserialize, Serialize, Debug, Clone, PartialEq, Eq)]
pub enum BalanceChanged {
    Increase(AssetAmount),
    Decrease(AssetAmount),
}

impl BalanceChanged {
    pub fn merge(&self, other: &Self) -> Self {
        match (self, other) {
            (BalanceChanged::Increase(a), BalanceChanged::Increase(b)) => {
                BalanceChanged::Increase(a.add(b))
            }
            (BalanceChanged::Decrease(a), BalanceChanged::Decrease(b)) => {
                BalanceChanged::Decrease(a.add(b))
            }
            (BalanceChanged::Increase(a), BalanceChanged::Decrease(b)) => {
                if a >= b {
                    BalanceChanged::Increase(a.sub(b))
                } else {
                    BalanceChanged::Decrease(b.sub(a))
                }
            }
            (BalanceChanged::Decrease(a), BalanceChanged::Increase(b)) => {
                if a >= b {
                    BalanceChanged::Decrease(a.sub(b))
                } else {
                    BalanceChanged::Increase(b.sub(a))
                }
            }
        }
    }
}

type WalletRegistrySeqId = VaultSeqId;
type BalanceVaultSeqId = VaultSeqId;
type SecretVaultSeqId = VaultSeqId;

#[derive(Debug, Clone)]
pub struct BlockWithShardingReports {
    pub block: Arc<Block>,
    pub consensus_report: CandidConsensusReport,
    pub core_registry_report: CandidCoreRegistryUpdateTxShardingSubReport,
    pub core_registry_data_integrity_report: DataIntegrityShardingSubReport,
    pub wallet_register_tx_reports:
        BTreeMap<WalletRegistrySeqId, Vec<CandidWalletRegisterTxShardingSubReport>>,
    pub wallet_registry_data_integrity_reports:
        BTreeMap<WalletRegistrySeqId, DataIntegrityShardingSubReport>,
    pub secret_update_tx_reports:
        BTreeMap<SecretVaultSeqId, Vec<CandidSecretUpdateTxShardingSubReport>>,
    pub secret_vault_data_integrity_reports:
        BTreeMap<SecretVaultSeqId, DataIntegrityShardingSubReport>,
    pub balance_changed_reports: BTreeMap<BalanceVaultSeqId, Vec<CandidBalanceShardingSubReport>>,
    pub withdrawal_reports: BTreeMap<BalanceVaultSeqId, Vec<CandidWithdrawalShardingSubReport>>,
    pub balance_vault_data_integrity_reports:
        BTreeMap<BalanceVaultSeqId, BalanceDataIntegrityShardingSubReport>,
}

impl BlockWithShardingReports {
    pub fn get_block(&self) -> Arc<Block> {
        self.block.clone()
    }

    pub fn get_consensus_report(&self) -> CandidConsensusReport {
        self.consensus_report.clone()
    }

    pub fn get_wallet_register_tx_reports(
        &self,
        wallet_vault_seq_id: &WalletRegistrySeqId,
    ) -> Vec<CandidWalletRegisterTxShardingSubReport> {
        self.wallet_register_tx_reports
            .get(wallet_vault_seq_id)
            .expect("Wallet register report should exist")
            .clone()
    }

    pub fn get_balance_changed_reports(
        &self,
        balance_vault_seq_id: &BalanceVaultSeqId,
    ) -> Vec<CandidBalanceShardingSubReport> {
        self.balance_changed_reports
            .get(balance_vault_seq_id)
            .expect("Balance changed report should exist")
            .clone()
    }

    pub fn get_withdrawal_reports(
        &self,
        balance_vault_seq_id: &BalanceVaultSeqId,
    ) -> Vec<CandidWithdrawalShardingSubReport> {
        self.withdrawal_reports
            .get(balance_vault_seq_id)
            .expect("Withdrawal report should exist")
            .clone()
    }

    pub fn get_secret_update_tx_reports(
        &self,
        secret_vault_seq_id: &SecretVaultSeqId,
    ) -> Vec<CandidSecretUpdateTxShardingSubReport> {
        self.secret_update_tx_reports
            .get(secret_vault_seq_id)
            .expect("Secret update tx report should exist")
            .clone()
    }

    pub fn get_data_integrity_reports_of_balance_vault(
        &self,
        balance_vault_seq_id: &BalanceVaultSeqId,
    ) -> BalanceDataIntegrityShardingSubReport {
        self.balance_vault_data_integrity_reports
            .get(balance_vault_seq_id)
            .expect("Data integrity report should exist")
            .clone()
    }

    pub fn get_data_integrity_reports_of_secret_vault(
        &self,
        secret_vault_seq_id: &SecretVaultSeqId,
    ) -> DataIntegrityShardingSubReport {
        self.secret_vault_data_integrity_reports
            .get(secret_vault_seq_id)
            .expect("Data integrity report should exist")
            .clone()
    }

    pub fn get_data_integrity_reports_of_wallet_registry(
        &self,
        wallet_registry_seq_id: &WalletRegistrySeqId,
    ) -> DataIntegrityShardingSubReport {
        self.wallet_registry_data_integrity_reports
            .get(wallet_registry_seq_id)
            .expect("Data integrity report should exist")
            .clone()
    }
}