use crate::{account_address::AccountAddress, on_chain_config::ValidatorSet};
use diem_crypto::{
    bls::deserialize_bls_public_key_unchecked, hash::CryptoHash, Signature,
    VRFPublicKey, VerifyingKey,
};
use serde::{Deserialize, Serialize};
use std::{collections::BTreeMap, fmt};
use thiserror::Error;
use crate::{
    on_chain_config::OnChainConfig,
    validator_config::{
        ConsensusPublicKey, ConsensusSignature, ConsensusVRFProof,
        ConsensusVRFPublicKey,
    },
};
#[cfg(any(test, feature = "fuzzing"))]
use anyhow::{ensure, Result};
#[cfg(any(test, feature = "fuzzing"))]
use proptest_derive::Arbitrary;
#[derive(Debug, Error, PartialEq)]
pub enum VerifyError {
    #[error("Author is unknown")]
    UnknownAuthor,
    #[error(
        "The voting power ({}) is less than quorum voting power ({})",
        voting_power,
        quorum_voting_power
    )]
    TooLittleVotingPower {
        voting_power: u64,
        quorum_voting_power: u64,
    },
    #[error(
        "The number of signatures ({}) is greater than total number of authors ({})",
        num_of_signatures,
        num_of_authors
    )]
    TooManySignatures {
        num_of_signatures: usize,
        num_of_authors: usize,
    },
    #[error("Signature is invalid")]
    InvalidSignature,
    #[error("Invalid VRF proof")]
    InvalidVrfProof,
}
#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[cfg_attr(any(test, feature = "fuzzing"), derive(Arbitrary))]
pub struct ValidatorConsensusInfo {
    #[serde(deserialize_with = "deserialize_bls_public_key_unchecked")]
    public_key: ConsensusPublicKey,
    vrf_public_key: Option<ConsensusVRFPublicKey>,
    voting_power: u64,
}
impl ValidatorConsensusInfo {
    pub fn new(
        public_key: ConsensusPublicKey,
        vrf_public_key: Option<ConsensusVRFPublicKey>, voting_power: u64,
    ) -> Self {
        ValidatorConsensusInfo {
            public_key,
            vrf_public_key,
            voting_power,
        }
    }
    pub fn voting_power(&self) -> u64 { self.voting_power }
    pub fn public_key(&self) -> &ConsensusPublicKey { &self.public_key }
    pub fn vrf_public_key(&self) -> &Option<ConsensusVRFPublicKey> {
        &self.vrf_public_key
    }
}
#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[cfg_attr(any(test, feature = "fuzzing"), derive(Arbitrary))]
pub struct ValidatorVerifier {
    address_to_validator_info: BTreeMap<AccountAddress, ValidatorConsensusInfo>,
    quorum_voting_power: u64,
    total_voting_power: u64,
}
impl ValidatorVerifier {
    pub fn new(
        address_to_validator_info: BTreeMap<
            AccountAddress,
            ValidatorConsensusInfo,
        >,
    ) -> Self {
        let total_voting_power = sum_voting_power(&address_to_validator_info);
        let quorum_voting_power = if address_to_validator_info.is_empty() {
            0
        } else {
            total_voting_power * 2 / 3 + 1
        };
        ValidatorVerifier {
            address_to_validator_info,
            quorum_voting_power,
            total_voting_power,
        }
    }
    #[cfg(any(test, feature = "fuzzing"))]
    pub fn new_with_quorum_voting_power(
        address_to_validator_info: BTreeMap<
            AccountAddress,
            ValidatorConsensusInfo,
        >,
        quorum_voting_power: u64,
    ) -> Result<Self> {
        let total_voting_power = sum_voting_power(&address_to_validator_info);
        ensure!(
            quorum_voting_power <= total_voting_power,
            "Quorum voting power is greater than the sum of all voting power of authors: {}, \
             quorum_size: {}.",
            quorum_voting_power,
            total_voting_power
        );
        Ok(ValidatorVerifier {
            address_to_validator_info,
            quorum_voting_power,
            total_voting_power,
        })
    }
    #[cfg(any(test, feature = "fuzzing"))]
    pub fn new_for_testing(
        address_to_validator_info: BTreeMap<
            AccountAddress,
            ValidatorConsensusInfo,
        >,
        quorum_voting_power: u64, total_voting_power: u64,
    ) -> Self {
        ValidatorVerifier {
            address_to_validator_info,
            quorum_voting_power,
            total_voting_power,
        }
    }
    pub fn new_single(
        author: AccountAddress, public_key: ConsensusPublicKey,
        vrf_public_key: Option<ConsensusVRFPublicKey>,
    ) -> Self {
        let mut author_to_validator_info = BTreeMap::new();
        author_to_validator_info.insert(
            author,
            ValidatorConsensusInfo::new(public_key, vrf_public_key, 1),
        );
        Self::new(author_to_validator_info)
    }
    pub fn verify<T: Serialize + CryptoHash>(
        &self, author: AccountAddress, message: &T,
        signature: &ConsensusSignature,
    ) -> std::result::Result<(), VerifyError> {
        match self.get_public_key(&author) {
            Some(public_key) => {
                if public_key
                    .verify_struct_signature(message, signature)
                    .is_err()
                {
                    Err(VerifyError::InvalidSignature)
                } else {
                    Ok(())
                }
            }
            None => Err(VerifyError::UnknownAuthor),
        }
    }
    pub fn verify_vrf(
        &self, author: AccountAddress, seed: &[u8], proof: &ConsensusVRFProof,
    ) -> std::result::Result<(), VerifyError> {
        match self.get_vrf_public_key(&author) {
            Some(Some(public_key)) => {
                if public_key.verify_proof(seed, proof).is_err() {
                    Err(VerifyError::InvalidVrfProof)
                } else {
                    Ok(())
                }
            }
            _ => Err(VerifyError::UnknownAuthor),
        }
    }
    pub fn verify_aggregated_struct_signature<T: CryptoHash + Serialize>(
        &self, message: &T,
        aggregated_signature: &BTreeMap<AccountAddress, ConsensusSignature>,
    ) -> std::result::Result<(), VerifyError> {
        self.check_num_of_signatures(aggregated_signature)?;
        self.check_voting_power(aggregated_signature.keys())?;
        for (author, signature) in aggregated_signature {
            self.verify(*author, message, &signature.clone())?;
        }
        Ok(())
    }
    pub fn batch_verify_aggregated_signatures<T: CryptoHash + Serialize>(
        &self, message: &T,
        aggregated_signature: &BTreeMap<AccountAddress, ConsensusSignature>,
    ) -> std::result::Result<(), VerifyError> {
        self.check_num_of_signatures(aggregated_signature)?;
        self.check_voting_power(aggregated_signature.keys())?;
        let keys_and_signatures: Vec<(ConsensusPublicKey, ConsensusSignature)> =
            aggregated_signature
                .iter()
                .flat_map(|(address, signature)| {
                    let sig = signature.clone();
                    self.get_public_key(&address).map(|pub_key| (pub_key, sig))
                })
                .collect();
        if ConsensusSignature::batch_verify(message, keys_and_signatures)
            .is_err()
        {
            self.verify_aggregated_struct_signature(
                message,
                aggregated_signature,
            )?
        }
        Ok(())
    }
    fn check_num_of_signatures(
        &self,
        aggregated_signature: &BTreeMap<AccountAddress, ConsensusSignature>,
    ) -> std::result::Result<(), VerifyError> {
        let num_of_signatures = aggregated_signature.len();
        if num_of_signatures > self.len() {
            return Err(VerifyError::TooManySignatures {
                num_of_signatures,
                num_of_authors: self.len(),
            });
        }
        Ok(())
    }
    pub fn check_voting_power<'a>(
        &self, authors: impl Iterator<Item = &'a AccountAddress>,
    ) -> std::result::Result<(), VerifyError> {
        let mut aggregated_voting_power = 0;
        for account_address in authors {
            match self.get_voting_power(&account_address) {
                Some(voting_power) => aggregated_voting_power += voting_power,
                None => return Err(VerifyError::UnknownAuthor),
            }
        }
        if aggregated_voting_power < self.quorum_voting_power {
            return Err(VerifyError::TooLittleVotingPower {
                voting_power: aggregated_voting_power,
                quorum_voting_power: self.quorum_voting_power,
            });
        }
        Ok(())
    }
    pub fn get_public_key(
        &self, author: &AccountAddress,
    ) -> Option<ConsensusPublicKey> {
        self.address_to_validator_info
            .get(&author)
            .map(|validator_info| validator_info.public_key.clone())
    }
    pub fn get_vrf_public_key(
        &self, author: &AccountAddress,
    ) -> Option<Option<ConsensusVRFPublicKey>> {
        self.address_to_validator_info
            .get(&author)
            .map(|validator_info| validator_info.vrf_public_key.clone())
    }
    pub fn get_voting_power(&self, author: &AccountAddress) -> Option<u64> {
        self.address_to_validator_info
            .get(&author)
            .map(|validator_info| validator_info.voting_power)
    }
    pub fn get_ordered_account_addresses_iter(
        &self,
    ) -> impl Iterator<Item = AccountAddress> + '_ {
        self.address_to_validator_info.keys().copied()
    }
    pub fn len(&self) -> usize { self.address_to_validator_info.len() }
    pub fn is_empty(&self) -> bool { self.len() == 0 }
    pub fn quorum_voting_power(&self) -> u64 { self.quorum_voting_power }
    pub fn total_voting_power(&self) -> u64 { self.total_voting_power }
    pub fn address_to_validator_info(
        &self,
    ) -> &BTreeMap<AccountAddress, ValidatorConsensusInfo> {
        &self.address_to_validator_info
    }
    pub fn extra_vote_count<'a>(
        &self, signers: impl Iterator<Item = &'a AccountAddress>,
    ) -> anyhow::Result<u64> {
        let mut total_count = 0;
        for signer in signers {
            total_count += self
                .get_voting_power(signer)
                .ok_or(anyhow::anyhow!("Signer is not a validator"))?;
        }
        total_count.checked_sub(self.quorum_voting_power).ok_or(
            anyhow::anyhow!("counted voting power overflows total power"),
        )
    }
}
impl OnChainConfig for ValidatorVerifier {
    const IDENTIFIER: &'static str = "DiemSystem";
}
fn sum_voting_power(
    address_to_validator_info: &BTreeMap<
        AccountAddress,
        ValidatorConsensusInfo,
    >,
) -> u64 {
    address_to_validator_info.values().fold(0, |sum, x| {
        sum.checked_add(x.voting_power)
            .expect("sum of all voting power is greater than u64::max")
    })
}
impl fmt::Display for ValidatorVerifier {
    fn fmt(&self, f: &mut fmt::Formatter) -> std::fmt::Result {
        write!(f, "ValidatorSet: [")?;
        for (addr, info) in &self.address_to_validator_info {
            write!(
                f,
                "{}: {}, ",
                addr.short_str_lossless(),
                info.voting_power
            )?;
        }
        write!(f, "]")
    }
}
impl From<&ValidatorSet> for ValidatorVerifier {
    fn from(validator_set: &ValidatorSet) -> Self {
        ValidatorVerifier::new(validator_set.payload().iter().fold(
            BTreeMap::new(),
            |mut map, validator| {
                map.insert(
                    *validator.account_address(),
                    ValidatorConsensusInfo::new(
                        validator.consensus_public_key().clone(),
                        validator.vrf_public_key().clone(),
                        validator.consensus_voting_power(),
                    ),
                );
                map
            },
        ))
    }
}
#[cfg(any(test, feature = "fuzzing"))]
impl From<&ValidatorVerifier> for ValidatorSet {
    fn from(verifier: &ValidatorVerifier) -> Self {
        ValidatorSet::new(
            verifier
                .get_ordered_account_addresses_iter()
                .map(|addr| {
                    crate::validator_info::ValidatorInfo::new_with_test_network_keys(
                        addr,
                        verifier.get_public_key(&addr).unwrap(),
                        verifier.get_vrf_public_key(&addr).unwrap(),
                        verifier.get_voting_power(&addr).unwrap(),
                    )
                })
                .collect(),
        )
    }
}
#[cfg(any(test, feature = "fuzzing"))]
pub fn random_validator_verifier(
    count: usize, custom_voting_power_quorum: Option<u64>,
    pseudo_random_account_address: bool,
) -> (
    Vec<crate::validator_signer::ValidatorSigner>,
    ValidatorVerifier,
) {
    let mut signers = Vec::new();
    let mut account_address_to_validator_info = BTreeMap::new();
    for i in 0..count {
        let random_signer = if pseudo_random_account_address {
            crate::validator_signer::ValidatorSigner::from_int(i as u8)
        } else {
            crate::validator_signer::ValidatorSigner::random([i as u8; 32])
        };
        account_address_to_validator_info.insert(
            random_signer.author(),
            crate::validator_verifier::ValidatorConsensusInfo::new(
                random_signer.public_key(),
                random_signer.vrf_public_key(),
                1,
            ),
        );
        signers.push(random_signer);
    }
    (
        signers,
        match custom_voting_power_quorum {
            Some(custom_voting_power_quorum) => {
                ValidatorVerifier::new_with_quorum_voting_power(
                    account_address_to_validator_info,
                    custom_voting_power_quorum,
                )
                .expect("Unable to create testing validator verifier")
            }
            None => ValidatorVerifier::new(account_address_to_validator_info),
        },
    )
}
#[cfg(test)]
mod tests {
    use super::*;
    use crate::validator_signer::ValidatorSigner;
    use diem_crypto::test_utils::{TestDiemCrypto, TEST_SEED};
    use std::collections::BTreeMap;
    #[test]
    fn test_check_voting_power() {
        let (validator_signers, validator_verifier) =
            random_validator_verifier(2, None, false);
        let mut author_to_signature_map = BTreeMap::new();
        assert_eq!(
            validator_verifier
                .check_voting_power(author_to_signature_map.keys())
                .unwrap_err(),
            VerifyError::TooLittleVotingPower {
                voting_power: 0,
                quorum_voting_power: 2,
            }
        );
        let dummy_struct = TestDiemCrypto("Hello, World".to_string());
        for validator in validator_signers.iter() {
            author_to_signature_map
                .insert(validator.author(), validator.sign(&dummy_struct));
        }
        assert_eq!(
            validator_verifier
                .check_voting_power(author_to_signature_map.keys()),
            Ok(())
        );
    }
    #[test]
    fn test_validator() {
        let validator_signer = ValidatorSigner::random(TEST_SEED);
        let dummy_struct = TestDiemCrypto("Hello, World".to_string());
        let signature = validator_signer.sign(&dummy_struct);
        let validator = ValidatorVerifier::new_single(
            validator_signer.author(),
            validator_signer.public_key(),
            validator_signer.vrf_public_key(),
        );
        assert_eq!(
            validator.verify(
                validator_signer.author(),
                &dummy_struct,
                &signature
            ),
            Ok(())
        );
        let unknown_validator_signer = ValidatorSigner::random([1; 32]);
        let unknown_signature = unknown_validator_signer.sign(&dummy_struct);
        assert_eq!(
            validator.verify(
                unknown_validator_signer.author(),
                &dummy_struct,
                &unknown_signature
            ),
            Err(VerifyError::UnknownAuthor)
        );
        assert_eq!(
            validator.verify(
                validator_signer.author(),
                &dummy_struct,
                &unknown_signature
            ),
            Err(VerifyError::InvalidSignature)
        );
    }
    #[test]
    fn test_equal_vote_quorum_validators() {
        const NUM_SIGNERS: u8 = 7;
        let validator_signers: Vec<ValidatorSigner> = (0..NUM_SIGNERS)
            .map(|i| ValidatorSigner::random([i; 32]))
            .collect();
        let dummy_struct = TestDiemCrypto("Hello, World".to_string());
        let mut author_to_public_key_map = BTreeMap::new();
        for validator in validator_signers.iter() {
            author_to_public_key_map.insert(
                validator.author(),
                ValidatorConsensusInfo::new(
                    validator.public_key(),
                    validator.vrf_public_key(),
                    1,
                ),
            );
        }
        let mut author_to_signature_map = BTreeMap::new();
        for validator in validator_signers.iter() {
            author_to_signature_map
                .insert(validator.author(), validator.sign(&dummy_struct));
        }
        let validator_verifier =
            ValidatorVerifier::new_with_quorum_voting_power(
                author_to_public_key_map,
                5,
            )
            .expect("Incorrect quorum size.");
        assert_eq!(
            validator_verifier.batch_verify_aggregated_signatures(
                &dummy_struct,
                &author_to_signature_map
            ),
            Ok(())
        );
        let unknown_validator_signer =
            ValidatorSigner::random([NUM_SIGNERS + 1; 32]);
        let unknown_signature = unknown_validator_signer.sign(&dummy_struct);
        author_to_signature_map.insert(
            unknown_validator_signer.author(),
            unknown_signature.clone(),
        );
        assert_eq!(
            validator_verifier.batch_verify_aggregated_signatures(
                &dummy_struct,
                &author_to_signature_map
            ),
            Err(VerifyError::TooManySignatures {
                num_of_signatures: 8,
                num_of_authors: 7
            })
        );
        author_to_signature_map.clear();
        for validator in validator_signers.iter().take(5) {
            author_to_signature_map
                .insert(validator.author(), validator.sign(&dummy_struct));
        }
        assert_eq!(
            validator_verifier.batch_verify_aggregated_signatures(
                &dummy_struct,
                &author_to_signature_map
            ),
            Ok(())
        );
        author_to_signature_map.insert(
            unknown_validator_signer.author(),
            unknown_signature.clone(),
        );
        assert_eq!(
            validator_verifier.batch_verify_aggregated_signatures(
                &dummy_struct,
                &author_to_signature_map
            ),
            Err(VerifyError::UnknownAuthor)
        );
        author_to_signature_map.clear();
        for validator in validator_signers.iter().take(4) {
            author_to_signature_map
                .insert(validator.author(), validator.sign(&dummy_struct));
        }
        assert_eq!(
            validator_verifier.batch_verify_aggregated_signatures(
                &dummy_struct,
                &author_to_signature_map
            ),
            Err(VerifyError::TooLittleVotingPower {
                voting_power: 4,
                quorum_voting_power: 5
            })
        );
        author_to_signature_map
            .insert(unknown_validator_signer.author(), unknown_signature);
        assert_eq!(
            validator_verifier.batch_verify_aggregated_signatures(
                &dummy_struct,
                &author_to_signature_map
            ),
            Err(VerifyError::UnknownAuthor)
        );
    }
    #[test]
    #[should_panic]
    fn test_very_unequal_vote_quorum_validators() {
        const NUM_SIGNERS: u8 = 4;
        let validator_signers: Vec<ValidatorSigner> = (0..NUM_SIGNERS)
            .map(|i| ValidatorSigner::random([i; 32]))
            .collect();
        let dummy_struct = TestDiemCrypto("Hello, World".to_string());
        let mut author_to_public_key_map = BTreeMap::new();
        let mut author_to_signature_map = BTreeMap::new();
        for (i, validator_signer) in validator_signers.iter().enumerate() {
            let mut voting_power: u64 = i as u64;
            if i == 3 {
                voting_power = u64::max_value()
            }
            author_to_public_key_map.insert(
                validator_signer.author(),
                ValidatorConsensusInfo::new(
                    validator_signer.public_key(),
                    validator_signer.vrf_public_key(),
                    voting_power,
                ),
            );
            author_to_signature_map.insert(
                validator_signer.author(),
                validator_signer.sign(&dummy_struct),
            );
        }
        let _validator_verifier =
            ValidatorVerifier::new(author_to_public_key_map);
    }
    #[test]
    fn test_unequal_vote_quorum_validators() {
        const NUM_SIGNERS: u8 = 4;
        let validator_signers: Vec<ValidatorSigner> = (0..NUM_SIGNERS)
            .map(|i| ValidatorSigner::random([i; 32]))
            .collect();
        let dummy_struct = TestDiemCrypto("Hello, World".to_string());
        let mut author_to_public_key_map = BTreeMap::new();
        let mut author_to_signature_map = BTreeMap::new();
        for (i, validator_signer) in validator_signers.iter().enumerate() {
            author_to_public_key_map.insert(
                validator_signer.author(),
                ValidatorConsensusInfo::new(
                    validator_signer.public_key(),
                    validator_signer.vrf_public_key(),
                    i as u64,
                ),
            );
            author_to_signature_map.insert(
                validator_signer.author(),
                validator_signer.sign(&dummy_struct),
            );
        }
        let validator_verifier =
            ValidatorVerifier::new_with_quorum_voting_power(
                author_to_public_key_map,
                5,
            )
            .expect("Incorrect quorum size.");
        assert_eq!(
            validator_verifier.batch_verify_aggregated_signatures(
                &dummy_struct,
                &author_to_signature_map
            ),
            Ok(())
        );
        let unknown_validator_signer =
            ValidatorSigner::random([NUM_SIGNERS + 1; 32]);
        let unknown_signature = unknown_validator_signer.sign(&dummy_struct);
        author_to_signature_map.insert(
            unknown_validator_signer.author(),
            unknown_signature.clone(),
        );
        assert_eq!(
            validator_verifier.batch_verify_aggregated_signatures(
                &dummy_struct,
                &author_to_signature_map
            ),
            Err(VerifyError::TooManySignatures {
                num_of_signatures: 5,
                num_of_authors: 4
            })
        );
        author_to_signature_map.clear();
        for validator in validator_signers.iter().skip(2) {
            author_to_signature_map
                .insert(validator.author(), validator.sign(&dummy_struct));
        }
        assert_eq!(
            validator_verifier.batch_verify_aggregated_signatures(
                &dummy_struct,
                &author_to_signature_map
            ),
            Ok(())
        );
        author_to_signature_map.insert(
            unknown_validator_signer.author(),
            unknown_signature.clone(),
        );
        assert_eq!(
            validator_verifier.batch_verify_aggregated_signatures(
                &dummy_struct,
                &author_to_signature_map
            ),
            Err(VerifyError::UnknownAuthor)
        );
        author_to_signature_map.clear();
        for validator in validator_signers.iter().take(3) {
            author_to_signature_map
                .insert(validator.author(), validator.sign(&dummy_struct));
        }
        assert_eq!(
            validator_verifier.batch_verify_aggregated_signatures(
                &dummy_struct,
                &author_to_signature_map
            ),
            Err(VerifyError::TooLittleVotingPower {
                voting_power: 3,
                quorum_voting_power: 5
            })
        );
        author_to_signature_map
            .insert(unknown_validator_signer.author(), unknown_signature);
        assert_eq!(
            validator_verifier.batch_verify_aggregated_signatures(
                &dummy_struct,
                &author_to_signature_map
            ),
            Err(VerifyError::UnknownAuthor)
        );
    }
}