// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

// Copyright 2021 Conflux Foundation. All rights reserved.
// Conflux is free software and distributed under GNU General Public License.
// See http://www.gnu.org/licenses/

//! PendingVotes store pending votes observed for a fixed epoch and round.
//! It is meant to be used inside of a RoundState.
//! The module takes care of creating a QC or a TC
//! when enough votes (or timeout votes) have been observed.
//! Votes are automatically dropped when the structure goes out of scope.

use consensus_types::{
    common::Author, quorum_cert::QuorumCert,
    timeout_certificate::TimeoutCertificate, vote::Vote,
};
use diem_crypto::{hash::CryptoHash, HashValue};
use diem_logger::prelude::*;
use diem_types::{
    ledger_info::LedgerInfoWithSignatures,
    validator_verifier::{ValidatorVerifier, VerifyError},
};
use std::{
    collections::{BTreeMap, HashMap},
    fmt,
    sync::Arc,
};

/// Result of the vote processing. The failure case (Verification error) is
/// returned as the Error part of the result.
#[derive(Debug, PartialEq)]
pub enum VoteReceptionResult {
    /// The vote has been added but QC has not been formed yet. Return the
    /// amount of voting power the given (proposal, execution) pair.
    VoteAdded(u64),
    /// The very same vote message has been processed in past.
    DuplicateVote,
    /// The very same author has already voted for another proposal in this
    /// round (equivocation).
    EquivocateVote((Vote, Vote)),
    /// This block has just been certified after adding the vote.
    NewQuorumCertificate(Arc<QuorumCert>),
    /// The vote completes a new TimeoutCertificate
    NewTimeoutCertificate(Arc<TimeoutCertificate>),
    /// There might be some issues adding a vote
    ErrorAddingVote(VerifyError),
    /// The vote is not for the current round.
    UnexpectedRound(u64, u64),
}

/// A PendingVotes structure keep track of votes
pub struct PendingVotes {
    /// Maps LedgerInfo digest to associated signatures (contained in a partial
    /// LedgerInfoWithSignatures). This might keep multiple LedgerInfos for
    /// the current round: either due to different proposals (byzantine
    /// behavior) or due to different NIL proposals (clients can have a
    /// different view of what block to extend).
    li_digest_to_votes: HashMap<
        HashValue, /* LedgerInfo digest */
        LedgerInfoWithSignatures,
    >,
    /// Tracks all the signatures of the votes for the given round. In case we
    /// succeed to aggregate 2f+1 signatures a TimeoutCertificate is
    /// formed.
    maybe_partial_tc: Option<TimeoutCertificate>,
    /// Map of Author to vote. This is useful to discard multiple votes.
    author_to_vote: HashMap<Author, Vote>,
}

impl PendingVotes {
    /// Creates an empty PendingVotes structure for a specific epoch and round
    pub fn new() -> Self {
        PendingVotes {
            li_digest_to_votes: HashMap::new(),
            maybe_partial_tc: None,
            author_to_vote: HashMap::new(),
        }
    }

    /// Insert a vote and if the vote is valid, return a QuorumCertificate
    /// preferentially over a TimeoutCertificate if either can can be formed
    pub fn insert_vote(
        &mut self, vote: &Vote, validator_verifier: &ValidatorVerifier,
    ) -> VoteReceptionResult {
        // derive data from vote
        let li_digest = vote.ledger_info().hash();

        //
        // 1. Has the author already voted for this round?
        //

        if let Some(previously_seen_vote) =
            self.author_to_vote.get(&vote.author())
        {
            // is it the same vote?
            if li_digest == previously_seen_vote.ledger_info().hash() {
                // we've already seen an equivalent vote before
                let new_timeout_vote =
                    vote.is_timeout() && !previously_seen_vote.is_timeout();
                if !new_timeout_vote {
                    // it's not a new timeout vote
                    return VoteReceptionResult::DuplicateVote;
                }
            } else {
                // we have seen a different vote for the same round
                diem_error!(
                    SecurityEvent::ConsensusEquivocatingVote,
                    remote_peer = vote.author(),
                    vote = vote,
                    previous_vote = previously_seen_vote
                );

                return VoteReceptionResult::EquivocateVote((
                    vote.clone(),
                    previously_seen_vote.clone(),
                ));
            }
        }

        //
        // 2. Store new vote (or update, in case it's a new timeout vote)
        //

        self.author_to_vote.insert(vote.author(), vote.clone());

        //
        // 3. Let's check if we can create a QC
        //

        // obtain the ledger info with signatures associated to the vote's
        // ledger info
        let li_with_sig =
            self.li_digest_to_votes.entry(li_digest).or_insert_with(|| {
                // if the ledger info with signatures doesn't exist yet, create
                // it
                LedgerInfoWithSignatures::new(
                    vote.ledger_info().clone(),
                    BTreeMap::new(),
                )
            });

        // add this vote to the ledger info with signatures
        li_with_sig.add_signature(vote.author(), vote.signature().clone());

        // check if we have enough signatures to create a QC
        let voting_power = match validator_verifier
            .check_voting_power(li_with_sig.signatures().keys())
        {
            // a quorum of signature was reached, a new QC is formed
            Ok(_) => {
                return VoteReceptionResult::NewQuorumCertificate(Arc::new(
                    QuorumCert::new(
                        vote.vote_data().clone(),
                        li_with_sig.clone(),
                    ),
                ));
            }

            // not enough votes
            Err(VerifyError::TooLittleVotingPower { voting_power, .. }) => {
                voting_power
            }

            // error
            Err(error) => {
                diem_error!(
                    "MUST_FIX: vote received could not be added: {}, vote: {}",
                    error,
                    vote
                );
                return VoteReceptionResult::ErrorAddingVote(error);
            }
        };

        //
        // 4. We couldn't form a QC, let's check if we can create a TC
        //

        if let Some(timeout_signature) = vote.timeout_signature() {
            // form timeout struct
            // TODO(mimoo): stronger: pass the (epoch, round) tuple as arguments
            // of this function
            let timeout = vote.timeout();

            // if no partial TC exist, create one
            let partial_tc = self
                .maybe_partial_tc
                .get_or_insert_with(|| TimeoutCertificate::new(timeout));

            // add the timeout signature
            partial_tc.add_signature(vote.author(), timeout_signature.clone());

            // did the TC reach a threshold?
            match validator_verifier
                .check_voting_power(partial_tc.signatures().keys())
            {
                // A quorum of signature was reached, a new TC was formed!
                Ok(_) => {
                    return VoteReceptionResult::NewTimeoutCertificate(
                        Arc::new(partial_tc.clone()),
                    )
                }

                // not enough votes
                Err(VerifyError::TooLittleVotingPower { .. }) => (),

                // error
                Err(error) => {
                    diem_error!(
                        "MUST_FIX: timeout vote received could not be added: {}, vote: {}",
                        error, vote
                    );
                    return VoteReceptionResult::ErrorAddingVote(error);
                }
            }
        }

        //
        // 5. No QC (or TC) could be formed, return the QC's voting power
        //

        VoteReceptionResult::VoteAdded(voting_power)
    }

    pub fn get_certificate(
        &self, validator_verifier: &ValidatorVerifier,
    ) -> VoteReceptionResult {
        for ledger_info in self.li_digest_to_votes.values() {
            if validator_verifier
                .check_voting_power(ledger_info.signatures().keys())
                .is_ok()
            {
                let vote = self
                    .author_to_vote
                    .get(
                        ledger_info
                            .signatures()
                            .keys()
                            .next()
                            .expect("enough voting power"),
                    )
                    .expect("vote was recorded");
                return VoteReceptionResult::NewQuorumCertificate(Arc::new(
                    QuorumCert::new(
                        vote.vote_data().clone(),
                        ledger_info.clone(),
                    ),
                ));
            }
        }
        if let Some(partial_tc) = &self.maybe_partial_tc {
            if validator_verifier
                .check_voting_power(partial_tc.signatures().keys())
                .is_ok()
            {
                return VoteReceptionResult::NewTimeoutCertificate(Arc::new(
                    partial_tc.clone(),
                ));
            }
        }
        // TODO(lpl): Use another result
        VoteReceptionResult::DuplicateVote
    }

    pub fn vote_received(&self, vote: &Vote) -> bool {
        let li_digest = vote.ledger_info().hash();

        if let Some(previously_seen_vote) =
            self.author_to_vote.get(&vote.author())
        {
            // is it the same vote?
            if li_digest == previously_seen_vote.ledger_info().hash() {
                // we've already seen an equivalent vote before
                let new_timeout_vote =
                    vote.is_timeout() && !previously_seen_vote.is_timeout();
                if !new_timeout_vote {
                    // it's not a new timeout vote
                    return true;
                }
            }
        }
        false
    }
}

//
// Helpful trait implementation
//

impl fmt::Display for PendingVotes {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        // collect votes per ledger info
        let votes = self
            .li_digest_to_votes
            .iter()
            .map(|(li_digest, li)| {
                (li_digest, li.signatures().keys().collect::<Vec<_>>())
            })
            .collect::<BTreeMap<_, _>>();

        // collect timeout votes
        let timeout_votes = self.maybe_partial_tc.as_ref().map(|partial_tc| {
            partial_tc.signatures().keys().collect::<Vec<_>>()
        });

        // write
        write!(f, "PendingVotes: [")?;

        for (hash, authors) in votes {
            write!(
                f,
                "LI {} has {} votes {:?} ",
                hash,
                authors.len(),
                authors
            )?;
        }

        if let Some(authors) = timeout_votes {
            write!(f, "{} timeout {:?}", authors.len(), authors)?;
        }

        write!(f, "]")
    }
}

//
// Tests
//
/*
#[cfg(test)]
mod tests {
    use super::{PendingVotes, VoteReceptionResult};
    use consensus_types::{vote::Vote, vote_data::VoteData};
    use diem_crypto::HashValue;
    use diem_types::{
        block_info::BlockInfo, ledger_info::LedgerInfo,
        validator_verifier::random_validator_verifier,
    };

    /// Creates a random ledger info for epoch 1 and round 1.
    fn random_ledger_info() -> LedgerInfo {
        LedgerInfo::new(
            BlockInfo::new(
                1,
                0,
                HashValue::random(),
                HashValue::random(),
                0,
                0,
                None,
            ),
            HashValue::random(),
        )
    }

    /// Creates a random VoteData for epoch 1 and round 1,
    /// extending a random block at epoch1 and round 0.
    fn random_vote_data() -> VoteData {
        VoteData::new(BlockInfo::random(1), BlockInfo::random(0))
    }

    #[test]
    /// Verify that votes are properly aggregated to QC based on their
    /// LedgerInfo digest
    fn test_qc_aggregation() {
        ::diem_logger::Logger::init_for_testing();

        // set up 4 validators
        let (signers, validator) = random_validator_verifier(4, Some(2), false);
        let mut pending_votes = PendingVotes::new();

        // create random vote from validator[0]
        let li1 = random_ledger_info();
        let vote_data_1 = random_vote_data();
        let vote_data_1_author_0 =
            Vote::new(vote_data_1, signers[0].author(), li1, &signers[0]);

        // first time a new vote is added -> VoteAdded
        assert_eq!(
            pending_votes.insert_vote(&vote_data_1_author_0, &validator),
            VoteReceptionResult::VoteAdded(1)
        );

        // same author voting for the same thing -> DuplicateVote
        assert_eq!(
            pending_votes.insert_vote(&vote_data_1_author_0, &validator),
            VoteReceptionResult::DuplicateVote
        );

        // same author voting for a different result -> EquivocateVote
        let li2 = random_ledger_info();
        let vote_data_2 = random_vote_data();
        let vote_data_2_author_0 = Vote::new(
            vote_data_2.clone(),
            signers[0].author(),
            li2.clone(),
            &signers[0],
        );
        assert_eq!(
            pending_votes.insert_vote(&vote_data_2_author_0, &validator),
            VoteReceptionResult::EquivocateVote
        );

        // a different author voting for a different result -> VoteAdded
        let vote_data_2_author_1 = Vote::new(
            vote_data_2.clone(),
            signers[1].author(),
            li2.clone(),
            &signers[1],
        );
        assert_eq!(
            pending_votes.insert_vote(&vote_data_2_author_1, &validator),
            VoteReceptionResult::VoteAdded(1)
        );

        // two votes for the ledger info -> NewQuorumCertificate
        let vote_data_2_author_2 =
            Vote::new(vote_data_2, signers[2].author(), li2, &signers[2]);
        match pending_votes.insert_vote(&vote_data_2_author_2, &validator) {
            VoteReceptionResult::NewQuorumCertificate(qc) => {
                assert!(validator
                    .check_voting_power(qc.ledger_info().signatures().keys())
                    .is_ok());
            }
            _ => {
                panic!("No QC formed.");
            }
        };
    }

    #[test]
    /// Verify that votes are properly aggregated to TC based on their rounds
    fn test_tc_aggregation() {
        ::diem_logger::Logger::init_for_testing();

        // set up 4 validators
        let (signers, validator) = random_validator_verifier(4, Some(2), false);
        let mut pending_votes = PendingVotes::new();

        // submit a new vote from validator[0] -> VoteAdded
        let li1 = random_ledger_info();
        let vote1 = random_vote_data();
        let mut vote1_author_0 =
            Vote::new(vote1, signers[0].author(), li1, &signers[0]);

        assert_eq!(
            pending_votes.insert_vote(&vote1_author_0, &validator),
            VoteReceptionResult::VoteAdded(1)
        );

        // submit the same vote but enhanced with a timeout -> VoteAdded
        let timeout = vote1_author_0.timeout();
        let signature = timeout.sign(&signers[0]);
        vote1_author_0.add_timeout_signature(signature);

        assert_eq!(
            pending_votes.insert_vote(&vote1_author_0, &validator),
            VoteReceptionResult::VoteAdded(1)
        );

        // another vote for a different block cannot form a TC if it doesn't
        // have a timeout signature
        let li2 = random_ledger_info();
        let vote2 = random_vote_data();
        let mut vote2_author_1 =
            Vote::new(vote2, signers[1].author(), li2, &signers[1]);
        assert_eq!(
            pending_votes.insert_vote(&vote2_author_1, &validator),
            VoteReceptionResult::VoteAdded(1)
        );

        // if that vote is now enhanced with a timeout signature ->
        // NewTimeoutCertificate
        let timeout = vote2_author_1.timeout();
        let signature = timeout.sign(&signers[1]);
        vote2_author_1.add_timeout_signature(signature);
        match pending_votes.insert_vote(&vote2_author_1, &validator) {
            VoteReceptionResult::NewTimeoutCertificate(tc) => {
                assert!(validator
                    .check_voting_power(tc.signatures().keys())
                    .is_ok());
            }
            _ => {
                panic!("No TC formed.");
            }
        };
    }
}
*/