use frost_core::{
    Identifier,
    keys::SigningShare,
    round1::{SigningCommitments, SigningNonces, commit},
};
use rand_core::CryptoRngCore;
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;
use zeroize::ZeroizeOnDrop;

use crate::{
    Ciphersuite, ReconstructionThreshold,
    errors::{InitializationError, ProtocolError},
    frost::sign_utils::assert_participant_inputs,
    participants::{Participant, ParticipantList},
    protocol::{
        Protocol,
        helpers::recv_from_others,
        internal::{Comms, SharedChannel, make_protocol},
    },
};

/// The necessary inputs for the creation of a presignature.
pub struct PresignArguments<C: Ciphersuite> {
    /// The output of key generation, i.e. our share of the secret key.
    pub private_share: SigningShare<C>,
    /// The threshold for the scheme
    pub threshold: ReconstructionThreshold,
}

/// The output of the presigning protocol.
///
/// This output is basically all the parts of the signature that we can perform
/// without knowing the message.
#[derive(Clone, Serialize, Deserialize, Eq, PartialEq, ZeroizeOnDrop)]
pub struct PresignOutput<C: Ciphersuite> {
    /// The secret signing nonces.
    pub nonces: SigningNonces<C>,
    #[zeroize(skip)]
    pub commitments_map: BTreeMap<Identifier<C>, SigningCommitments<C>>,
}

impl_secret_debug!({C: Ciphersuite} PresignOutput<C> { show: [commitments_map], redact: [nonces] });

/// Maximum incoming buffer entries for the FROST presign protocol.
const FROST_PRESIGN_MAX_INCOMING_BUFFER_ENTRIES: usize = 1;

/// Runs Presigning of either `EdDSA` or `RedDSA`
pub fn presign<C, R>(
    participants: &[Participant],
    me: Participant,
    args: &PresignArguments<C>,
    rng: R,
) -> Result<impl Protocol<Output = PresignOutput<C>> + use<C, R>, InitializationError>
where
    C: Ciphersuite,
    R: CryptoRngCore + Send + 'static,
{
    let participants = assert_participant_inputs(participants, args.threshold, me)?;

    let ctx = Comms::with_buffer_capacity(FROST_PRESIGN_MAX_INCOMING_BUFFER_ENTRIES);
    let fut = do_presign(
        ctx.shared_channel(),
        participants,
        me,
        args.private_share,
        rng,
    );
    Ok(make_protocol(ctx, fut))
}

async fn do_presign<C: Ciphersuite>(
    mut chan: SharedChannel,
    participants: ParticipantList,
    me: Participant,
    signing_share: SigningShare<C>,
    mut rng: impl CryptoRngCore,
) -> Result<PresignOutput<C>, ProtocolError> {
    // --- Round 1
    // * Compute and Receive commitments.
    let mut commitments_map: BTreeMap<Identifier<C>, SigningCommitments<C>> = BTreeMap::new();

    // Step 1.1
    // Creating two commitments and corresponding nonces
    let (nonces, commitments) = commit(&signing_share, &mut rng);
    commitments_map.insert(me.to_identifier()?, commitments);

    // Step 1.2
    let commit_waitpoint = chan.next_waitpoint();
    // Sending the commitments to all
    chan.send_many(commit_waitpoint, &commitments)?;

    // Step 1.3 and 1.4
    // Collecting the commitments
    for (from, commitment) in recv_from_others(&chan, commit_waitpoint, &participants, me).await? {
        commitments_map.insert(from.to_identifier()?, commitment);
    }

    Ok(PresignOutput {
        nonces,
        commitments_map,
    })
}

#[cfg(test)]
mod test {
    use super::*;
    use crate::test_utils::{MockCryptoRng, assert_buffer_capacity, generate_participants};
    use frost_ed25519::Ed25519Sha512;
    use rand::SeedableRng;
    use rstest::rstest;

    #[rstest]
    #[case(3, 2)]
    #[case(5, 3)]
    #[case(10, 4)]
    fn test_presign_buffer_entries(#[case] num_participants: usize, #[case] threshold: usize) {
        // Given
        let participants = generate_participants(num_participants);
        let mut rng = MockCryptoRng::seed_from_u64(42);
        let keygen_result = crate::test_utils::run_keygen::<Ed25519Sha512, MockCryptoRng>(
            &participants,
            threshold,
            &mut rng,
        );

        // When + Then
        assert_buffer_capacity(
            &participants,
            &mut rng,
            |comms, p_list, p, rng_p| {
                let private_share = keygen_result
                    .iter()
                    .find(|(kp, _)| *kp == p)
                    .unwrap()
                    .1
                    .private_share;
                do_presign::<Ed25519Sha512>(comms.shared_channel(), p_list, p, private_share, rng_p)
            },
            |_| FROST_PRESIGN_MAX_INCOMING_BUFFER_ENTRIES,
        );
    }
}