use std::collections::HashMap;

use rand::Rng;
use rand_core::{CryptoRngCore, SeedableRng};

use threshold_signatures::{
    ReconstructionThreshold,
    frost::eddsa,
    participants::Participant,
    protocol::Protocol,
    test_utils::{MockCryptoRng, generate_participants_with_random_ids, run_keygen},
};

use super::{MAX_MALICIOUS, PreparedPresig};

pub type PresignProtocols = Vec<(
    Participant,
    Box<dyn Protocol<Output = eddsa::PresignOutput>>,
)>;

pub fn ed25519_build_presign_protocols<R: CryptoRngCore + SeedableRng + Send + 'static>(
    participants: &[Participant],
    key_packages: &[(Participant, eddsa::KeygenOutput)],
    threshold: ReconstructionThreshold,
    rng: &mut R,
) -> (PresignProtocols, HashMap<Participant, u64>) {
    let mut protocols = Vec::with_capacity(participants.len());
    let mut seeds = HashMap::with_capacity(participants.len());
    for (p, keygen_out) in key_packages {
        let seed = rng.next_u64();
        let rng_p = MockCryptoRng::seed_from_u64(seed);
        let protocol = eddsa::presign(
            participants,
            *p,
            &eddsa::PresignArguments {
                private_share: keygen_out.private_share,
                threshold,
            },
            rng_p,
        )
        .map(|presig| Box::new(presig) as Box<dyn Protocol<Output = eddsa::PresignOutput>>)
        .expect("Presignature should succeed");
        protocols.push((*p, protocol));
        seeds.insert(*p, seed);
    }
    (protocols, seeds)
}

/// Used to prepare ed25519 presignatures for benchmarking
pub fn ed25519_prepare_presign<R: CryptoRngCore + SeedableRng + Send + 'static>(
    num_participants: usize,
    rng: &mut R,
) -> FrostEd25519PreparedPresig {
    let participants = generate_participants_with_random_ids(num_participants, rng);
    let key_packages = run_keygen(&participants, *MAX_MALICIOUS + 1, rng);
    let threshold = ReconstructionThreshold::from(*MAX_MALICIOUS + 1);
    let (protocols, seeds) =
        ed25519_build_presign_protocols(&participants, &key_packages, threshold, rng);
    FrostEd25519PreparedPresig {
        protocols,
        key_packages,
        participants,
        seeds,
    }
}

/// Used to prepare ed25519 signatures for benchmarking
pub fn ed25519_prepare_sign_v1<R: CryptoRngCore + SeedableRng + Send + 'static>(
    threshold: ReconstructionThreshold,
    rng: &mut R,
) -> FrostEd25519SigV1 {
    let num_participants = threshold.value();
    let participants = generate_participants_with_random_ids(num_participants, rng);
    let key_packages = run_keygen(&participants, *MAX_MALICIOUS + 1, rng);

    // choose a coordinator at random
    let coordinator_index = rng.gen_range(0..num_participants);
    let coordinator = participants[coordinator_index];

    let mut protocols = Vec::with_capacity(participants.len());
    let mut seeds = HashMap::with_capacity(participants.len());

    let mut message: [u8; 32] = [0u8; 32];
    rng.fill_bytes(&mut message);
    let message = message.to_vec();

    for (p, keygen_out) in &key_packages {
        let seed = rng.next_u64();
        let rng_p = MockCryptoRng::seed_from_u64(seed);
        let protocol = eddsa::sign::sign_v1(
            &participants,
            threshold,
            *p,
            coordinator,
            keygen_out.clone(),
            message.clone(),
            rng_p,
        )
        .map(|sig| Box::new(sig) as Box<dyn Protocol<Output = eddsa::SignatureOption>>)
        .expect("Signing should succeed");
        protocols.push((*p, protocol));
        seeds.insert(*p, seed);
    }

    FrostEd25519SigV1 {
        protocols,
        index: coordinator_index,
        key_packages,
        message,
        seeds,
    }
}

pub fn ed25519_prepare_sign_v2<R: CryptoRngCore + SeedableRng + Send + 'static>(
    result: &[(Participant, eddsa::PresignOutput)],
    key_packages: Vec<(Participant, eddsa::KeygenOutput)>,
    threshold: ReconstructionThreshold,
    rng: &mut R,
) -> FrostEd25519SigV2 {
    let num_participants = threshold.value();
    // collect all participants
    let participants: Vec<_> = result.iter().map(|(participant, _)| *participant).collect();

    // choose a coordinator at random
    let coordinator_index = rng.gen_range(0..num_participants);
    let coordinator = participants[coordinator_index];

    let mut protocols = Vec::with_capacity(participants.len());

    let mut message: [u8; 32] = [0u8; 32];
    rng.fill_bytes(&mut message);
    let message = message.to_vec();

    for ((p, keygen_out), (p_redundancy, presign)) in key_packages.iter().zip(result) {
        assert_eq!(p, p_redundancy);
        let protocol = eddsa::sign::sign_v2(
            &participants,
            threshold,
            *p,
            coordinator,
            keygen_out.clone(),
            presign.clone(),
            message.clone(),
        )
        .map(|sig| Box::new(sig) as Box<dyn Protocol<Output = eddsa::SignatureOption>>)
        .expect("Signing should succeed");
        protocols.push((*p, protocol));
    }

    FrostEd25519SigV2 {
        protocols,
        index: coordinator_index,
        presig: result[coordinator_index].1.clone(),
        key_packages,
        message,
    }
}

pub struct FrostEd25519SigV1 {
    pub protocols: Vec<(
        Participant,
        Box<dyn Protocol<Output = eddsa::SignatureOption>>,
    )>,
    pub index: usize,
    pub key_packages: Vec<(Participant, eddsa::KeygenOutput)>,
    pub message: Vec<u8>,
    pub seeds: HashMap<Participant, u64>,
}

pub type FrostEd25519PreparedPresig = PreparedPresig<eddsa::PresignOutput, eddsa::KeygenOutput>;
pub struct FrostEd25519SigV2 {
    pub protocols: Vec<(
        Participant,
        Box<dyn Protocol<Output = eddsa::SignatureOption>>,
    )>,
    pub index: usize,
    pub presig: eddsa::PresignOutput,
    pub key_packages: Vec<(Participant, eddsa::KeygenOutput)>,
    pub message: Vec<u8>,
}