use candid::{CandidType, Deserialize, Nat};
use ex3_canister_constant::{START_ASSET_ID, START_SPOT_MARKET_ID, START_WALLET_ID};
use ex3_node_types::block::{BlockHeader, CandidBlockHeader, CandidConsensusBlockHeader};
use ex3_node_types::number::Ex3Uint;
use ex3_node_types::{
    AssetId, CandidAssetId, CandidSpotMarketId, CandidWalletRegisterId, CanisterId, SpotMarketId,
    WalletRegisterId,
};
use ex3_serde::bincode;
use ic_stable_structures::storable::Bound;
use ic_stable_structures::Storable;
use num_traits::{CheckedAdd, CheckedSub, One, Zero};
use serde::Serialize;
use serde_bytes::ByteBuf;

pub mod error;

type BlockchainArchiveCanister = CanisterId;

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

#[derive(CandidType, Deserialize, Debug, Clone)]
pub struct CandidConsensusReport {
    pub head: CandidBlockHeader,
    #[serde(with = "serde_bytes")]
    pub cid: Vec<u8>,
    pub activated_vault_report: NewSnapshotVaultActivatedReport,
    pub id_cursor_report: IdCursorReport,
    pub withdrawal_counter_report: WithdrawalCounterReport,
}

#[derive(Deserialize, Serialize, Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
pub struct ConsensusReport {
    pub head: BlockHeader,
    #[serde(with = "serde_bytes")]
    pub cid: Vec<u8>,
    pub activated_vault_report: NewSnapshotVaultActivatedReport,
    pub id_cursor_report: IdCursorReport,
    pub withdrawal_counter_report: WithdrawalCounterReport,
}

impl Storable for ConsensusReport {
    fn to_bytes(&self) -> std::borrow::Cow<[u8]> {
        bincode::serialize(&self).unwrap().into()
    }

    fn from_bytes(bytes: std::borrow::Cow<[u8]>) -> Self {
        bincode::deserialize(&bytes).unwrap()
    }

    const BOUND: Bound = Bound::Bounded {
        max_size: 1056,
        is_fixed_size: false,
    };
}

#[derive(
    CandidType, Deserialize, Serialize, Default, Debug, Clone, PartialEq, Eq, PartialOrd, Ord,
)]
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, PartialOrd, Ord,
)]
pub struct IdCursorReport {
    pub wallet_id_cursor_increase: u32,
    pub asset_id_cursor_increase: u32,
    pub market_id_cursor_increase: u32,
}

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

impl Storable for WithdrawalCounterReport {
    fn to_bytes(&self) -> std::borrow::Cow<[u8]> {
        let mut bytes = self.accepted_withdrawal_count.to_le_bytes().to_vec();
        bytes.extend_from_slice(&self.rejected_withdrawal_count.to_le_bytes());
        bytes.into()
    }

    fn from_bytes(bytes: std::borrow::Cow<[u8]>) -> Self {
        let mut bytes = bytes.into_owned();
        let accepted_withdrawal_count =
            u32::from_le_bytes(bytes.drain(..4).collect::<Vec<u8>>().try_into().unwrap());
        let rejected_withdrawal_count =
            u32::from_le_bytes(bytes.drain(..4).collect::<Vec<u8>>().try_into().unwrap());
        Self {
            accepted_withdrawal_count,
            rejected_withdrawal_count,
        }
    }

    const BOUND: Bound = Bound::Bounded {
        max_size: 8,
        is_fixed_size: true,
    };
}

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

#[derive(CandidType, Deserialize, Serialize, Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
pub struct ActivatedSnapshotVault {
    pub canister_seq_id: u64,
    pub canister_total_count: u32,
    pub canister_index: u32,
    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 spot market id
    pub max_spot_market_id: Option<CandidSpotMarketId>,
}

#[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: Ex3Uint,

    /// 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 spot market id
    pub max_spot_market_id: Option<SpotMarketId>,
}

impl Storable for SnapshotCursor {
    fn to_bytes(&self) -> std::borrow::Cow<[u8]> {
        bincode::serialize(&self).unwrap().into()
    }

    fn from_bytes(bytes: std::borrow::Cow<[u8]>) -> Self {
        bincode::deserialize(&bytes).unwrap()
    }

    const BOUND: Bound = Bound::Bounded {
        max_size: 155,
        is_fixed_size: false,
    };
}

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 += Ex3Uint::from(activated_vault_report.wallet_registries.len());
        self.vault_count += Ex3Uint::from(activated_vault_report.secret_vaults.len());
        self.vault_count += Ex3Uint::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(
            WalletRegisterId::from(wallet_id_cursor_increase.clone())
                .checked_add(&START_WALLET_ID.into())
                .unwrap()
                .checked_sub(&Ex3Uint::one())
                .expect("Invalid max wallet id: overflow"),
            |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(
            AssetId::from(asset_id_cursor_increase.clone())
                .checked_add(&START_ASSET_ID.into())
                .unwrap()
                .checked_sub(&Ex3Uint::one())
                .expect("Invalid max asset id: overflow"),
            |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_spot_market_id.clone();
        let new_max_market_id = current_max_market_id.map_or(
            SpotMarketId::from(market_id_cursor_increase.clone())
                .checked_add(&START_SPOT_MARKET_ID.into())
                .unwrap()
                .checked_sub(&Ex3Uint::one())
                .expect("Invalid max market id: overflow"),
            |id| id + market_id_cursor_increase,
        );

        self.max_spot_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: Ex3Uint::zero(),
            max_wallet_id: None,
            max_asset_id: None,
            max_spot_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.into()),
            max_asset_id: candid_snapshot_cursor
                .max_asset_id
                .map(|asset_id| asset_id.into()),
            max_spot_market_id: candid_snapshot_cursor
                .max_spot_market_id
                .map(|market_id| market_id.into()),
        }
    }
}

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_spot_market_id: snapshot_cursor.max_spot_market_id.map(Into::into),
        }
    }
}

impl From<CandidConsensusReport> for ConsensusReport {
    fn from(candid_consensus_report: CandidConsensusReport) -> Self {
        Self {
            head: candid_consensus_report.head.into(),
            cid: candid_consensus_report.cid,
            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,
        }
    }
}

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

    #[test]
    fn test_storable_for_consensus_report() {
        let consensus_report = ConsensusReport {
            head: BlockHeader {
                version: 2,
                height: u128::MAX.into(),
                p_id: u128::MAX.into(),
                pre_block_hash: [0; 32],
                tx_root: [1; 32],
                state_changed_root: [2; 32],
                data_integrity_root: [3; 32],
                rejected_tx_root: [4; 32],
            },
            cid: ByteBuf::from(vec![0; 32]),
            activated_vault_report: NewSnapshotVaultActivatedReport {
                wallet_registries: vec![ActivatedWalletRegistry {
                    canister_seq_id: 0,
                    canister_total_count: 0,
                    canister_index: 0,
                    hash_of_sharding: [0; 32],
                }],
                secret_vaults: vec![ActivatedSecretVault {
                    canister_seq_id: 0,
                    canister_total_count: 0,
                    canister_index: 0,
                    hash_of_sharding: [0; 32],
                }],
                balance_vaults: vec![ActivatedBalanceVault {
                    canister_seq_id: 0,
                    canister_total_count: 0,
                    canister_index: 0,
                    hash_of_sharding: [0; 32],
                }],
            },
            id_cursor_report: IdCursorReport::default(),
            withdrawal_counter_report: WithdrawalCounterReport::default(),
        };
        let bytes = consensus_report.to_bytes();
        assert!(bytes.len() <= 1056, "bytes.len() = {}", bytes.len());
        let consensus_report2 = ConsensusReport::from_bytes(bytes);
        assert_eq!(consensus_report, consensus_report2);
    }
}