//! This module and the frost one are supposed to have the same helper function
use super::{KeygenOutput, PresignOutput, SignatureOption};
use crate::{
    ReconstructionThreshold,
    crypto::random::Randomness,
    errors::{InitializationError, ProtocolError},
    frost::assert_sign_inputs,
    participants::{Participant, ParticipantList},
    protocol::{
        Protocol,
        helpers::recv_from_others,
        internal::{Comms, SharedChannel, make_protocol},
    },
};

use reddsa::frost::redjubjub::{
    CheaterDetection, Identifier, SigningPackage,
    keys::{KeyPackage, PublicKeyPackage},
    rerandomized,
    rerandomized::RandomizedParams,
    round2::SignatureShare,
};
use std::collections::BTreeMap;
use zeroize::Zeroizing;

/// Depending on whether the current participant is a coordinator or not,
/// runs the signature protocol as either a participant or a coordinator.
///
/// WARNING: Extracted from FROST documentation:
/// In all of the main FROST ciphersuites, the entire message must be sent
/// to participants. In some cases, where the message is too big, it may be
/// necessary to send a hash of the message instead. We strongly suggest
/// creating a specific ciphersuite for this, and not just sending the hash
/// as if it were the message.
/// For reference, see how RFC 8032 handles "pre-hashing".
///
/// Maximum incoming buffer entries for the coordinator in the `RedJubjub` sign protocol.
pub(crate) const REDJUBJUB_SIGN_MAX_INCOMING_COORDINATOR_ENTRIES: usize = 1;
/// Maximum incoming buffer entries for non-coordinator participants in the `RedJubjub` sign protocol.
#[cfg(test)]
pub(crate) const REDJUBJUB_SIGN_MAX_INCOMING_PARTICIPANT_ENTRIES: usize = 1;

/// /!\ Warning: the threshold in this scheme is the exactly the
///              same as the max number of malicious parties.
///
/// The `randomizer_seed` is the upstream-recommended way to randomize a FROST
/// signing run: the coordinator passes `Some(seed)` and other participants pass
/// `None`. The seed is forwarded to participants, who deterministically derive
/// the randomizer from `(seed, signing_commitments)` so they don't have to
/// trust the coordinator's RNG.
#[allow(clippy::too_many_arguments)]
pub fn sign<T>(
    participants: &[Participant],
    threshold: T,
    me: Participant,
    coordinator: Participant,
    keygen_output: KeygenOutput,
    presignature: PresignOutput,
    message: Vec<u8>,
    randomizer_seed: Option<Randomness>,
) -> Result<impl Protocol<Output = SignatureOption> + use<T>, InitializationError>
where
    T: Into<ReconstructionThreshold>,
{
    let threshold = threshold.into();
    let participants = assert_sign_inputs(participants, threshold, me, coordinator)?;

    let comms = Comms::with_buffer_capacity(REDJUBJUB_SIGN_MAX_INCOMING_COORDINATOR_ENTRIES);
    let chan = comms.shared_channel();
    let fut = fut_wrapper(
        chan,
        participants,
        threshold,
        me,
        coordinator,
        keygen_output,
        presignature,
        message,
        randomizer_seed,
    );
    Ok(make_protocol(comms, fut))
}

#[allow(clippy::too_many_arguments)]
async fn fut_wrapper(
    chan: SharedChannel,
    participants: ParticipantList,
    threshold: ReconstructionThreshold,
    me: Participant,
    coordinator: Participant,
    keygen_output: KeygenOutput,
    presignature: PresignOutput,
    message: Vec<u8>,
    randomizer_seed: Option<Randomness>,
) -> Result<SignatureOption, ProtocolError> {
    if me == coordinator {
        match randomizer_seed {
            Some(seed) => {
                do_sign_coordinator(
                    chan,
                    participants,
                    threshold,
                    me,
                    keygen_output,
                    &presignature,
                    message,
                    seed,
                )
                .await
            }
            None => Err(ProtocolError::InvalidInput(
                "Randomizer seed should not be none for coordinator".to_string(),
            )),
        }
    } else {
        match randomizer_seed {
            Some(_) => Err(ProtocolError::InvalidInput(
                "Randomizer seed should be none for non-coordinator".to_string(),
            )),
            None => {
                do_sign_participant(
                    chan,
                    threshold,
                    me,
                    coordinator,
                    keygen_output,
                    &presignature,
                    message,
                )
                .await
            }
        }
    }
}

/// Returns a future that executes signature protocol for *the Coordinator*.
///
/// WARNING: Extracted from FROST documentation:
/// In all of the main FROST ciphersuites, the entire message must be sent
/// to participants. In some cases, where the message is too big, it may be
/// necessary to send a hash of the message instead. We strongly suggest
/// creating a specific ciphersuite for this, and not just sending the hash
/// as if it were the message.
/// For reference, see how RFC 8032 handles "pre-hashing".
#[allow(clippy::too_many_arguments)]
async fn do_sign_coordinator(
    mut chan: SharedChannel,
    participants: ParticipantList,
    threshold: ReconstructionThreshold,
    me: Participant,
    keygen_output: KeygenOutput,
    presignature: &PresignOutput,
    message: Vec<u8>,
    randomizer_seed: Randomness,
) -> Result<SignatureOption, ProtocolError> {
    // --- Round 1
    let key_package = construct_key_package(threshold, me, &keygen_output)?;
    let key_package = Zeroizing::new(key_package);
    let signing_package = SigningPackage::new(presignature.commitments_map.clone(), &message);
    let randomized_params = RandomizedParams::regenerate_from_seed_and_commitments(
        &keygen_output.public_key,
        randomizer_seed.as_ref(),
        &presignature.commitments_map,
    )
    .map_err(|_| ProtocolError::ErrorFrostSigningFailed)?;

    // Send the randomizer seed to everyone; participants regenerate the
    // randomizer themselves from the seed and signing commitments.
    let wait_round_1 = chan.next_waitpoint();
    chan.send_many(wait_round_1, &randomizer_seed)?;

    // Round 2
    let signature_share = rerandomized::sign_with_randomizer_seed(
        &signing_package,
        &presignature.nonces,
        &key_package,
        randomizer_seed.as_ref(),
    )
    .map_err(|_| ProtocolError::ErrorFrostSigningFailed)?;

    let sign_waitpoint = chan.next_waitpoint();
    let mut signature_shares: BTreeMap<Identifier, SignatureShare> = BTreeMap::new();
    signature_shares.insert(me.to_identifier()?, signature_share);
    for (from, signature_share) in
        recv_from_others(&chan, sign_waitpoint, &participants, me).await?
    {
        signature_shares.insert(from.to_identifier()?, signature_share);
    }

    // --- Signature aggregation.
    // * Converted collected signature shares into the signature.
    // * Signature is verified internally during `aggregate()` call.

    // We use empty BTreeMap because "cheater-detection" feature is disabled
    // Feature "cheater-detection" unveils existant malicious participants
    let pk_package = PublicKeyPackage::new(BTreeMap::new(), keygen_output.public_key, None);

    // Use `Disabled` cheater detection because we intentionally pass an empty
    // verifying-shares map in `pk_package`. The default `aggregate` (which uses
    // `FirstCheater`) would reject this with `UnknownIdentifier`.
    let signature = rerandomized::aggregate_custom(
        &signing_package,
        &signature_shares,
        &pk_package,
        CheaterDetection::Disabled,
        &randomized_params,
    )
    .map_err(|_| ProtocolError::ErrorFrostAggregation)?;
    Ok(Some(signature))
}

/// Returns a future that executes signature protocol for *a Participant*.
///
/// WARNING: Extracted from FROST documentation:
/// In all of the main FROST ciphersuites, the entire message must be sent
/// to participants. In some cases, where the message is too big, it may be
/// necessary to send a hash of the message instead. We strongly suggest
/// creating a specific ciphersuite for this, and not just sending the hash
/// as if it were the message.
/// For reference, see how RFC 8032 handles "pre-hashing".
async fn do_sign_participant(
    mut chan: SharedChannel,
    threshold: ReconstructionThreshold,
    me: Participant,
    coordinator: Participant,
    keygen_output: KeygenOutput,
    presignature: &PresignOutput,
    message: Vec<u8>,
) -> Result<SignatureOption, ProtocolError> {
    // --- Round 1.
    if coordinator == me {
        return Err(ProtocolError::InvalidInput(
            "the do_sign_participant function cannot be called
            for a coordinator"
                .to_string(),
        ));
    }

    // Receive the randomizer seed from the coordinator
    let wait_round_1 = chan.next_waitpoint();
    let randomizer_seed = loop {
        let (from, seed): (_, Randomness) = chan.recv(wait_round_1).await?;
        if from != coordinator {
            continue;
        }
        break seed;
    };

    let key_package = construct_key_package(threshold, me, &keygen_output)?;
    let key_package = Zeroizing::new(key_package);
    let signing_package = SigningPackage::new(presignature.commitments_map.clone(), &message);
    // nonces are zeroized when presignature drops (ZeroizeOnDrop)
    let signature_share = rerandomized::sign_with_randomizer_seed(
        &signing_package,
        &presignature.nonces,
        &key_package,
        randomizer_seed.as_ref(),
    )
    .map_err(|_| ProtocolError::ErrorFrostSigningFailed)?;

    let sign_waitpoint = chan.next_waitpoint();
    chan.send_private(sign_waitpoint, coordinator, &signature_share)?;

    Ok(None)
}

/// A function that takes a signing share and a keygenOutput
/// and construct a public key package used for frost signing
fn construct_key_package(
    threshold: ReconstructionThreshold,
    me: Participant,
    keygen_output: &KeygenOutput,
) -> Result<KeyPackage, ProtocolError> {
    let identifier = me.to_identifier()?;
    let signing_share = keygen_output.private_share;
    let verifying_share = signing_share.into();
    let verifying_key = keygen_output.public_key;
    let key_package = KeyPackage::new(
        identifier,
        signing_share,
        verifying_share,
        verifying_key,
        u16::try_from(threshold.value()).map_err(|_| {
            ProtocolError::Other("threshold cannot be converted to u16".to_string())
        })?,
    );

    // Ensures the values are zeroized on drop
    Ok(key_package)
}

#[cfg(test)]
mod test {
    use super::{
        REDJUBJUB_SIGN_MAX_INCOMING_COORDINATOR_ENTRIES,
        REDJUBJUB_SIGN_MAX_INCOMING_PARTICIPANT_ENTRIES, fut_wrapper,
    };
    use crate::test_utils::{
        assert_buffer_capacity, build_frost_key_packages_with_dealer, expected_buffer_by_role,
    };
    use crate::{
        Protocol,
        crypto::{hash::hash, random::Randomness},
        frost::redjubjub::{
            PresignOutput, SignatureOption, sign::sign, test::run_sign_with_presign,
        },
        test_utils::{MockCryptoRng, check_one_coordinator_output},
    };
    use frost_core::Field;
    use rand::{SeedableRng, seq::SliceRandom as _};
    use rand_core::RngCore;
    use reddsa::frost::redjubjub::{JubjubBlake2b512, JubjubScalarField, round1::commit};
    use rstest::rstest;
    use std::collections::BTreeMap;

    #[test]
    fn stress() {
        let mut rng = MockCryptoRng::seed_from_u64(42);

        let max_signers = 7;
        let msg = "hello_near";
        let msg_hash = hash(&msg).unwrap();

        for threshold in 2..max_signers {
            for actual_signers in threshold..=max_signers {
                let key_packages =
                    build_frost_key_packages_with_dealer(max_signers, threshold, &mut rng);
                let threshold: usize = threshold.into();
                let coordinator = key_packages[0].0;
                let data = run_sign_with_presign(
                    &key_packages,
                    actual_signers.into(),
                    coordinator,
                    threshold,
                    msg_hash,
                )
                .unwrap();
                check_one_coordinator_output(data, coordinator).unwrap();
            }
        }
    }

    #[test]
    fn test_signature_correctness() {
        let mut rng = MockCryptoRng::seed_from_u64(42);
        let threshold = 6;
        let keys = build_frost_key_packages_with_dealer(11, threshold, &mut rng);
        let public_key = keys[0].1.public_key.to_element();

        let msg = b"hello world".to_vec();
        let coordinator = keys.choose(&mut rng).expect("keys list is not empty").0;
        let mut participants_sign_builder = keys
            .iter()
            .map(|(p, keygen_output)| {
                let rng_p = MockCryptoRng::seed_from_u64(rng.next_u64());
                (*p, (keygen_output.clone(), rng_p))
            })
            .collect::<Vec<_>>();

        let mut commitments_map = BTreeMap::new();
        let mut nonces_map = BTreeMap::new();
        for (p, (keygen, rng_p)) in &mut participants_sign_builder {
            // Creating two commitments and corresponding nonces
            let (nonces, commitments) = commit(&keygen.private_share, rng_p);
            commitments_map.insert(p.to_identifier().unwrap(), commitments);
            nonces_map.insert(*p, nonces);
        }

        let mut rng = MockCryptoRng::seed_from_u64(644_221);
        let randomizer_seed = Randomness::random(&mut rng);
        // This checks the output signature validity internally
        let result = crate::test_utils::run_sign::<JubjubBlake2b512, _, _, _>(
            participants_sign_builder,
            coordinator,
            public_key,
            JubjubScalarField::zero(), // not important
            |participants, coordinator, me, _, (keygen_output, _), _| {
                let nonces = nonces_map.get(&me).unwrap().clone();
                let presignature = PresignOutput {
                    nonces,
                    commitments_map: commitments_map.clone(),
                };
                let seed = if me == coordinator {
                    Some(randomizer_seed.clone())
                } else {
                    None
                };
                sign(
                    participants,
                    threshold as usize,
                    me,
                    coordinator,
                    keygen_output,
                    presignature,
                    msg.clone(),
                    seed,
                )
                .map(|sig| Box::new(sig) as Box<dyn Protocol<Output = SignatureOption>>)
            },
        )
        .unwrap();
        let signature = check_one_coordinator_output(result, coordinator).unwrap();
        insta::assert_json_snapshot!(signature);
    }

    #[rstest]
    #[case(3, 2)]
    #[case(5, 3)]
    #[case(10, 4)]
    fn test_sign_buffer_entries(#[case] num_participants: usize, #[case] threshold: usize) {
        // Given
        let mut rng = MockCryptoRng::seed_from_u64(42);
        let keys = build_frost_key_packages_with_dealer(
            u16::try_from(num_participants).unwrap(),
            u16::try_from(threshold).unwrap(),
            &mut rng,
        );
        let coordinator = keys[0].0;
        let message = b"test message for buffer entries".to_vec();

        let mut presig_rng = MockCryptoRng::seed_from_u64(100);
        let mut commitments_map = std::collections::BTreeMap::new();
        let mut nonces_map = std::collections::HashMap::new();
        for (p, keygen) in &keys {
            let (nonces, commitments) =
                reddsa::frost::redjubjub::round1::commit(&keygen.private_share, &mut presig_rng);
            commitments_map.insert(p.to_identifier().unwrap(), commitments);
            nonces_map.insert(*p, nonces);
        }

        let randomizer_seed = Randomness::random(&mut presig_rng);

        let participants_list: Vec<_> = keys.iter().map(|(p, _)| *p).collect();

        // When + Then
        assert_buffer_capacity(
            &participants_list,
            &mut rng,
            |comms, p_list, p, _rng_p| {
                let keygen_output = keys.iter().find(|(kp, _)| *kp == p).unwrap().1.clone();
                let presign_output = PresignOutput {
                    nonces: nonces_map[&p].clone(),
                    commitments_map: commitments_map.clone(),
                };
                let seed = if p == coordinator {
                    Some(randomizer_seed.clone())
                } else {
                    None
                };
                fut_wrapper(
                    comms.shared_channel(),
                    p_list,
                    threshold.into(),
                    p,
                    coordinator,
                    keygen_output,
                    presign_output,
                    message.clone(),
                    seed,
                )
            },
            expected_buffer_by_role(
                coordinator,
                REDJUBJUB_SIGN_MAX_INCOMING_COORDINATOR_ENTRIES,
                REDJUBJUB_SIGN_MAX_INCOMING_PARTICIPANT_ENTRIES,
            ),
        );
    }
}