use {
    crate::common::{certificate_limits_and_vote_types, VoteType},
    agave_votor_messages::consensus_message::{Certificate, CertificateType, VoteMessage},
    bitvec::prelude::*,
    solana_bls_signatures::{BlsError, SignatureProjective},
    solana_signer_store::{encode_base2, encode_base3, EncodeError},
    thiserror::Error,
};

/// Maximum number of validators in a certificate
///
/// There are around 1500 validators currently. For a clean power-of-two
/// implementation, we should choose either 2048 or 4096. Choose a more
/// conservative number 4096 for now. During build() we will cut off end
/// of the bitmaps if the tail contains only zeroes, so actual bitmap
/// length will be less than or equal to this number.
const MAXIMUM_VALIDATORS: usize = 4096;

/// Different types of errors that can be returned from the [`CertificateBuilder::aggregate()`] function.
#[derive(Debug, PartialEq, Eq, Error)]
pub(super) enum AggregateError {
    #[error("BLS error: {0}")]
    Bls(#[from] BlsError),
    #[error("Invalid rank: {0}")]
    InvalidRank(u16),
    #[error("Validator already included")]
    ValidatorAlreadyIncluded,
}

/// Different types of errors that can be returned from the [`CertificateBuilder::build()`] function.
#[derive(Debug, PartialEq, Eq, Error)]
pub enum BuildError {
    #[error("BLS error: {0}")]
    Bls(#[from] BlsError),
    #[error("Encoding failed: {0:?}")]
    Encode(EncodeError),
    #[error("Invalid rank: {0}")]
    InvalidRank(usize),
}

/// Looks up the bit at `rank` in `bitmap` and sets it to true.
fn try_set_bitmap(bitmap: &mut BitVec<u8, Lsb0>, rank: u16) -> Result<(), AggregateError> {
    let mut ptr = bitmap
        .get_mut(rank as usize)
        .ok_or(AggregateError::InvalidRank(rank))?;
    if *ptr {
        return Err(AggregateError::ValidatorAlreadyIncluded);
    }
    *ptr = true;
    Ok(())
}

/// A builder for creating a `CertificateMessage` by efficiently aggregating BLS signatures.
#[derive(Clone)]
pub struct CertificateBuilder {
    cert_type: CertificateType,
    signature: SignatureProjective,
    // For some certificates we need two bitmaps, for example, NotarizeFallback
    // certificates have Notarize and NotarizeFallback votes, so we need two bitmaps
    // to represent them. The order of the VoteType is defined in certificate_limits_and_vote_types.
    // We normally put fallback votes in the second bitmap.
    // The order of the VoteType is important, if you change it, you might interpret
    // the bitmap incorrectly.
    // Some certificates (like Finalize) only need one bitmap, then the second bitmap
    // will be empty.
    input_bitmap_1: BitVec<u8, Lsb0>,
    input_bitmap_2: BitVec<u8, Lsb0>,
}

impl CertificateBuilder {
    pub fn new(cert_type: CertificateType) -> Self {
        Self {
            cert_type,
            signature: SignatureProjective::identity(),
            input_bitmap_1: BitVec::repeat(false, MAXIMUM_VALIDATORS),
            input_bitmap_2: BitVec::repeat(false, MAXIMUM_VALIDATORS),
        }
    }

    /// Aggregates a slice of `VoteMessage`s into the builder.
    pub fn aggregate(&mut self, messages: &[VoteMessage]) -> Result<(), AggregateError> {
        let vote_types = certificate_limits_and_vote_types(&self.cert_type).1;
        for vote_message in messages {
            let rank = vote_message.rank;

            let current_vote_type = VoteType::get_type(&vote_message.vote);

            if current_vote_type == vote_types[0] {
                try_set_bitmap(&mut self.input_bitmap_1, rank)?;
            } else if vote_types.len() == 2 && current_vote_type == vote_types[1] {
                try_set_bitmap(&mut self.input_bitmap_2, rank)?;
            }
        }

        Ok(self
            .signature
            .aggregate_with(messages.iter().map(|m| &m.signature))?)
    }

    pub fn build(self) -> Result<Certificate, BuildError> {
        let mut input_bitmap_1 = self.input_bitmap_1;
        let mut input_bitmap_2 = self.input_bitmap_2;

        let last_one_1 = input_bitmap_1 // use local variable
            .last_one()
            .map_or(0, |i| i.saturating_add(1));
        let last_one_2 = input_bitmap_2 // use local variable
            .last_one()
            .map_or(0, |i| i.saturating_add(1));
        let new_length = last_one_1.max(last_one_2);
        // checks in `aggregate()` guarantee that this assertion is valid
        debug_assert!(new_length <= MAXIMUM_VALIDATORS);
        if new_length > MAXIMUM_VALIDATORS {
            return Err(BuildError::InvalidRank(new_length));
        }

        input_bitmap_1.resize(new_length, false);
        input_bitmap_2.resize(new_length, false);
        let bitmap = if input_bitmap_2.count_ones() > 0 {
            // If we have two bitmaps, use Base3 encoding
            encode_base3(&input_bitmap_1, &input_bitmap_2).map_err(BuildError::Encode)?
        } else {
            // If we only have one bitmap, use Base2 encoding
            encode_base2(&input_bitmap_1).map_err(BuildError::Encode)?
        };
        Ok(Certificate {
            cert_type: self.cert_type,
            signature: self.signature.into(),
            bitmap,
        })
    }
}

#[cfg(test)]
mod tests {
    use {
        super::*,
        agave_votor_messages::{
            consensus_message::{CertificateType, VoteMessage},
            vote::Vote,
        },
        solana_bls_signatures::{
            Keypair as BLSKeypair, PubkeyProjective as BLSPubkeyProjective,
            Signature as BLSSignature, SignatureProjective, VerifiablePubkey,
        },
        solana_hash::Hash,
        solana_signer_store::{decode, Decoded},
    };

    #[test]
    fn test_normal_build() {
        let hash = Hash::new_unique();
        let cert_type = CertificateType::NotarizeFallback(1, hash);
        let mut builder = CertificateBuilder::new(cert_type);
        // Test building the certificate from Notarize and NotarizeFallback votes
        // Create Notarize on validator 1, 4, 6
        let vote = Vote::new_notarization_vote(1, hash);
        let rank_1 = [1, 4, 6];
        let messages_1 = rank_1
            .iter()
            .map(|&rank| {
                let keypair = BLSKeypair::new();
                let signature = keypair.sign(b"fake_vote_message");
                VoteMessage {
                    vote,
                    signature: signature.into(),
                    rank,
                }
            })
            .collect::<Vec<_>>();
        builder
            .aggregate(&messages_1)
            .expect("Failed to aggregate notarization votes");
        // Create NotarizeFallback on validator 2, 3, 5, 7
        let vote = Vote::new_notarization_fallback_vote(1, hash);
        let rank_2 = [2, 3, 5, 7];
        let messages_2 = rank_2
            .iter()
            .map(|&rank| {
                let keypair = BLSKeypair::new();
                let signature = keypair.sign(b"fake_vote_message_2");
                VoteMessage {
                    vote,
                    signature: signature.into(),
                    rank,
                }
            })
            .collect::<Vec<_>>();
        builder
            .aggregate(&messages_2)
            .expect("Failed to aggregate notarization fallback votes");

        let cert = builder.build().expect("Failed to build certificate");
        assert_eq!(cert.cert_type, cert_type);
        match decode(&cert.bitmap, MAXIMUM_VALIDATORS).expect("Failed to decode bitmap") {
            Decoded::Base3(bitmap1, bitmap2) => {
                assert_eq!(bitmap1.len(), 8);
                assert_eq!(bitmap2.len(), 8);
                for i in rank_1 {
                    assert!(bitmap1[i as usize]);
                }
                assert_eq!(bitmap1.count_ones(), 3);
                for i in rank_2 {
                    assert!(bitmap2[i as usize]);
                }
                assert_eq!(bitmap2.count_ones(), 4);
            }
            _ => panic!("Expected Base3 encoding"),
        }

        // Build a new certificate with only Notarize votes, we should get Base2 encoding
        let mut builder = CertificateBuilder::new(cert_type);
        builder
            .aggregate(&messages_1)
            .expect("Failed to aggregate notarization votes");
        let cert = builder.build().expect("Failed to build certificate");
        assert_eq!(cert.cert_type, cert_type);
        match decode(&cert.bitmap, MAXIMUM_VALIDATORS).expect("Failed to decode bitmap") {
            Decoded::Base2(bitmap1) => {
                assert_eq!(bitmap1.len(), 7);
                for i in rank_1 {
                    assert!(bitmap1[i as usize]);
                }
                assert_eq!(bitmap1.count_ones(), 3);
            }
            _ => panic!("Expected Base2 encoding"),
        }

        // Base2 encoding only applies when the first bitmap is non-empty, if we build another
        // certificate with only NotarizeFallback votes, we should still get Base3 encoding
        let mut builder = CertificateBuilder::new(cert_type);
        builder
            .aggregate(&messages_2)
            .expect("Failed to aggregate notarization fallback votes");
        let cert = builder.build().expect("Failed to build certificate");
        assert_eq!(cert.cert_type, cert_type);
        match decode(&cert.bitmap, MAXIMUM_VALIDATORS).expect("Failed to decode bitmap") {
            Decoded::Base3(bitmap1, bitmap2) => {
                assert_eq!(bitmap1.count_ones(), 0);
                assert_eq!(bitmap2.len(), 8);
                for i in rank_2 {
                    assert!(bitmap2[i as usize]);
                }
                assert_eq!(bitmap2.count_ones(), 4);
            }
            _ => panic!("Expected Base3 encoding"),
        }
    }

    #[test]
    fn test_builder_with_errors() {
        let hash = Hash::new_unique();
        let cert_type = CertificateType::NotarizeFallback(1, hash);
        let mut builder = CertificateBuilder::new(cert_type);

        // Test with a rank that exceeds the maximum allowed
        let vote = Vote::new_notarization_vote(1, hash);
        let vote2 = Vote::new_notarization_fallback_vote(1, hash);
        let rank_out_of_bounds = MAXIMUM_VALIDATORS.saturating_add(1); // Exceeds MAXIMUM_VALIDATORS
        let keypair = BLSKeypair::new();
        let signature = keypair.sign(b"fake_vote_message");
        let message_out_of_bounds = VoteMessage {
            vote,
            signature: signature.into(),
            rank: rank_out_of_bounds as u16,
        };
        assert_eq!(
            builder.aggregate(&[message_out_of_bounds]),
            Err(AggregateError::InvalidRank(rank_out_of_bounds as u16))
        );

        // Test bls error
        let message_with_invalid_signature = VoteMessage {
            vote,
            signature: BLSSignature::default(), // Invalid signature
            rank: 1,
        };
        assert_eq!(
            builder.aggregate(&[message_with_invalid_signature]),
            Err(AggregateError::Bls(BlsError::PointConversion))
        );

        // Test encoding error
        // Create two bitmaps with the same rank set
        let signature = keypair.sign(b"fake_vote_message_2");
        let messages_1 = vec![VoteMessage {
            vote,
            signature: signature.into(),
            rank: 1,
        }];
        let mut builder = CertificateBuilder::new(cert_type);
        builder
            .aggregate(&messages_1)
            .expect("Failed to aggregate notarization votes");
        let messages_2 = vec![VoteMessage {
            vote: vote2,
            signature: signature.into(),
            rank: 1, // Same rank as in messages_1
        }];
        builder
            .aggregate(&messages_2)
            .expect("Failed to aggregate notarization fallback votes");
        assert_eq!(
            builder.build(),
            Err(BuildError::Encode(EncodeError::InvalidBitCombination))
        );
    }

    #[test]
    fn test_certificate_verification_base2_encoding() {
        let slot = 10;
        let hash = Hash::new_unique();
        let cert_type = CertificateType::Notarize(slot, hash);

        // 1. Setup: Create keypairs and a single vote object.
        // All validators will sign the same message, resulting in a single bitmap.
        let num_validators = 5;
        let mut keypairs = Vec::new();
        let mut vote_messages = Vec::new();
        let vote = Vote::new_notarization_vote(slot, hash);
        let serialized_vote = bincode::serialize(&vote).unwrap();

        for i in 0..num_validators {
            let keypair = BLSKeypair::new();
            let signature = keypair.sign(&serialized_vote);
            vote_messages.push(VoteMessage {
                vote,
                signature: signature.into(),
                rank: i as u16,
            });
            keypairs.push(keypair);
        }

        // 2. Generation: Aggregate votes and build the certificate. This will
        // use base2 encoding because it only contains one type of vote.
        let mut builder = CertificateBuilder::new(cert_type);
        builder
            .aggregate(&vote_messages)
            .expect("Failed to aggregate votes");
        let certificate_message = builder.build().expect("Failed to build certificate");

        // 3. Verification: Aggregate the public keys and verify the signature.
        let aggregate_pubkey = BLSPubkeyProjective::aggregate(keypairs.iter().map(|kp| &kp.public))
            .expect("Failed to aggregate public keys");

        let verification_result =
            aggregate_pubkey.verify_signature(&certificate_message.signature, &serialized_vote);

        assert!(
            verification_result.unwrap_or(false),
            "BLS aggregate signature verification failed for base2 encoded certificate"
        );
    }

    #[test]
    fn test_certificate_verification_base3_encoding() {
        let slot = 20;
        let hash = Hash::new_unique();
        // A NotarizeFallback certificate can be composed of both Notarize and NotarizeFallback
        // votes.
        let cert_type = CertificateType::NotarizeFallback(slot, hash);

        // 1. Setup: Create two groups of validators signing two different vote types.
        let mut all_vote_messages = Vec::new();
        let mut all_pubkeys = Vec::new();
        let mut all_messages = Vec::new();

        // Group 1: Signs a Notarize vote.
        let notarize_vote = Vote::new_notarization_vote(slot, hash);
        let serialized_notarize_vote = bincode::serialize(&notarize_vote).unwrap();
        for i in 0..3 {
            let keypair = BLSKeypair::new();
            let signature = keypair.sign(&serialized_notarize_vote);
            all_vote_messages.push(VoteMessage {
                vote: notarize_vote,
                signature: signature.into(),
                rank: i as u16, // Ranks 0, 1, 2
            });
            all_pubkeys.push(keypair.public);
            all_messages.push(serialized_notarize_vote.clone());
        }

        // Group 2: Signs a NotarizeFallback vote.
        let notarize_fallback_vote = Vote::new_notarization_fallback_vote(slot, hash);
        let serialized_fallback_vote = bincode::serialize(&notarize_fallback_vote).unwrap();
        for i in 3..6 {
            let keypair = BLSKeypair::new();
            let signature = keypair.sign(&serialized_fallback_vote);
            all_vote_messages.push(VoteMessage {
                vote: notarize_fallback_vote,
                signature: signature.into(),
                rank: i as u16, // Ranks 3, 4, 5
            });
            all_pubkeys.push(keypair.public);
            all_messages.push(serialized_fallback_vote.clone());
        }

        // 2. Generation: Aggregate votes. Because there are two vote types, this will use
        //    base3 encoding.
        let mut builder = CertificateBuilder::new(cert_type);
        builder
            .aggregate(&all_vote_messages)
            .expect("Failed to aggregate votes");
        let certificate_message = builder.build().expect("Failed to build certificate");

        // 3. Verification:
        let decoded_bitmap =
            decode(&certificate_message.bitmap, MAXIMUM_VALIDATORS).expect("Failed to decode");

        match decoded_bitmap {
            Decoded::Base2(_bitmap) => {
                panic!("Expected Base3 encoding, but got Base2 encoding");
            }
            Decoded::Base3(bitmap1, bitmap2) => {
                // Bitmap1 should correspond to the Notarize votes (ranks 0, 1, 2)
                assert_eq!(bitmap1.count_ones(), 3);
                assert!(bitmap1[0] && bitmap1[1] && bitmap1[2]);
                // Bitmap2 should correspond to the NotarizeFallback votes (ranks 3, 4, 5)
                assert_eq!(bitmap2.count_ones(), 3);
                assert!(bitmap2[3] && bitmap2[4] && bitmap2[5]);
            }
        }

        SignatureProjective::verify_distinct_aggregated(
            all_pubkeys.iter(),
            &certificate_message.signature,
            all_messages.iter().map(Vec::as_slice),
        )
        .unwrap();
    }
}
