use elliptic_curve::scalar::IsHigh;

use crate::{
    MaxMalicious,
    ecdsa::{
        AffinePoint, Scalar, Secp256K1Sha256, Signature, SignatureOption,
        robust_ecdsa::RerandomizedPresignOutput, x_coordinate,
    },
    errors::{InitializationError, ProtocolError},
    participants::{Participant, ParticipantList},
    protocol::{
        Protocol,
        helpers::recv_from_others,
        internal::{Comms, SharedChannel, make_protocol},
    },
};
use frost_core::serialization::SerializableScalar;
use subtle::ConditionallySelectable;
type C = Secp256K1Sha256;

/// Maximum incoming buffer entries for the coordinator in the robust ECDSA sign protocol.
pub(crate) const ROBUST_ECDSA_SIGN_MAX_INCOMING_COORDINATOR_ENTRIES: usize = 1;
/// Maximum incoming buffer entries for non-coordinator participants in the robust ECDSA sign protocol.
#[cfg(test)]
pub(crate) const ROBUST_ECDSA_SIGN_MAX_INCOMING_PARTICIPANT_ENTRIES: usize = 0;

/// Depending on whether the current participant is a coordinator or not,
/// runs the signature protocol as either a participant or a coordinator.
///
/// WARNING:
/// This robust ECDSA scheme is vulnerable to split-view attacks in the robust
/// setting if different subsets of participants sign different `(msg_hash, tweak)`
/// values using shares derived from the same presignature (i.e., different
/// rerandomization inputs for the same presignature).
/// To reduce risk in this implementation, require `N1 = N2 = 2 * max_malicious + 1`,
/// ensure all participants agree on `(msg_hash, tweak, participants)` when creating
/// `RerandomizedPresignOutput`, never reuse a presignature, and do not sign with
/// `msg_hash == 0`.
pub fn sign<M>(
    participants: &[Participant],
    coordinator: Participant,
    max_malicious: M,
    me: Participant,
    public_key: AffinePoint,
    presignature: RerandomizedPresignOutput,
    msg_hash: Scalar,
) -> Result<impl Protocol<Output = SignatureOption> + use<M>, InitializationError>
where
    M: Into<MaxMalicious>,
{
    if participants.len() < 2 {
        return Err(InitializationError::NotEnoughParticipants {
            participants: participants.len(),
        });
    }

    let participants =
        ParticipantList::new(participants).ok_or(InitializationError::DuplicateParticipants)?;

    // ensure my presence in the participant list
    if !participants.contains(me) {
        return Err(InitializationError::MissingParticipant {
            role: "self",
            participant: me,
        });
    }

    // ensure the coordinator is a participant
    if !participants.contains(coordinator) {
        return Err(InitializationError::MissingParticipant {
            role: "coordinator",
            participant: coordinator,
        });
    }

    // ensure number of participants during the signing phase is >= 2 * max_malicious + 1
    let robust_ecdsa_threshold = max_malicious
        .into()
        .value()
        .checked_mul(2)
        .and_then(|v| v.checked_add(1))
        .ok_or_else(|| {
            InitializationError::BadParameters(
                "2*threshold+1 must be less than usize::MAX".to_string(),
            )
        })?;
    if robust_ecdsa_threshold > participants.len() {
        return Err(InitializationError::BadParameters(
            "2*max_malicious+1 must be less than or equals to participant count".to_string(),
        ));
    }

    // The next two conditions prevent split-view attacks
    // documented in docs/ecdsa/robust_ecdsa/signing.md
    if participants.len() != robust_ecdsa_threshold {
        return Err(InitializationError::BadParameters(
            "the number of participants during signing must be exactly 2*max_malicious+1 to avoid split view attacks".to_string(),
        ));
    }
    if bool::from(msg_hash.is_zero()) {
        return Err(InitializationError::BadParameters(
            "msg_hash cannot be 0 to avoid potential split view attacks".to_string(),
        ));
    }

    let ctx = Comms::with_buffer_capacity(ROBUST_ECDSA_SIGN_MAX_INCOMING_COORDINATOR_ENTRIES);
    let fut = fut_wrapper(
        ctx.shared_channel(),
        participants,
        coordinator,
        me,
        public_key,
        presignature,
        msg_hash,
    );
    Ok(make_protocol(ctx, fut))
}

/// Performs signing from any participant's perspective (except the coordinator)
fn do_sign_participant(
    mut chan: SharedChannel,
    participants: &ParticipantList,
    coordinator: Participant,
    me: Participant,
    presignature: &RerandomizedPresignOutput,
    msg_hash: Scalar,
) -> Result<SignatureOption, ProtocolError> {
    let s_me = compute_signature_share(presignature, msg_hash, participants, me)?;
    let wait_round = chan.next_waitpoint();
    chan.send_private(wait_round, coordinator, &s_me)?;

    Ok(None)
}

/// Performs signing from only the coordinator's perspective
async fn do_sign_coordinator(
    mut chan: SharedChannel,
    participants: ParticipantList,
    me: Participant,
    public_key: AffinePoint,
    presignature: RerandomizedPresignOutput,
    msg_hash: Scalar,
) -> Result<SignatureOption, ProtocolError> {
    let mut s = compute_signature_share(&presignature, msg_hash, &participants, me)?.0;
    let wait_round = chan.next_waitpoint();

    for (_, s_i) in
        recv_from_others::<SerializableScalar<C>>(&chan, wait_round, &participants, me).await?
    {
        // Sum the linearized shares
        s += s_i.0;
    }

    // raise error if s is zero
    if s.is_zero().into() {
        return Err(ProtocolError::AssertionFailed(
            "signature part s cannot be zero".to_string(),
        ));
    }
    // Normalize s
    s.conditional_assign(&(-s), s.is_high());

    let sig = Signature {
        big_r: presignature.big_r,
        s,
    };

    if !sig.verify(&public_key, &msg_hash) {
        return Err(ProtocolError::AssertionFailed(
            "signature failed to verify".to_string(),
        ));
    }

    Ok(Some(sig))
}

/// A common computation done by both the coordinator and the other participants
fn compute_signature_share(
    presignature: &RerandomizedPresignOutput,
    msg_hash: Scalar,
    participants: &ParticipantList,
    me: Participant,
) -> Result<SerializableScalar<C>, ProtocolError> {
    // (beta_i + tweak * k_i) * delta^{-1}
    let big_r = presignature.big_r;
    let big_r_x_coordinate = x_coordinate(&big_r);
    // beta * Rx + e
    let beta = presignature.beta * big_r_x_coordinate + presignature.e;

    let s_me = msg_hash * presignature.alpha + beta;
    // lambda_i * s_i
    let linearized_s_me = s_me * participants.lagrange::<C>(me)?;
    Ok(SerializableScalar::<C>(linearized_s_me))
}

/// Wraps the coordinator and the participant into a single functions to be called
async fn fut_wrapper(
    chan: SharedChannel,
    participants: ParticipantList,
    coordinator: Participant,
    me: Participant,
    public_key: AffinePoint,
    presignature: RerandomizedPresignOutput,
    msg_hash: Scalar,
) -> Result<SignatureOption, ProtocolError> {
    if me == coordinator {
        do_sign_coordinator(chan, participants, me, public_key, presignature, msg_hash).await
    } else {
        do_sign_participant(
            chan,
            &participants,
            coordinator,
            me,
            &presignature,
            msg_hash,
        )
    }
}

#[cfg(test)]
mod test {

    use k256::{PublicKey, ecdsa::VerifyingKey, ecdsa::signature::Verifier};
    use rand_core::{CryptoRngCore, SeedableRng};

    use super::*;
    use crate::ecdsa::{
        Field, Polynomial, ProjectivePoint, Secp256K1ScalarField,
        robust_ecdsa::{
            PresignOutput,
            test::{run_sign_with_rerandomization, run_sign_without_rerandomization},
        },
    };
    use crate::test_utils::{
        MockCryptoRng, assert_buffer_capacity, expected_buffer_by_role, generate_participants,
    };
    use rstest::rstest;

    type PresigSimulationOutput = (Scalar, Polynomial, Polynomial, Polynomial, ProjectivePoint);

    fn simulate_presignature(
        max_malicious: impl Into<MaxMalicious>,
        rng: &mut impl CryptoRngCore,
    ) -> PresigSimulationOutput {
        let max_malicious = max_malicious.into().value();
        // the presignatures scheme requires the generation of 5 different polynomials
        // (fk, fa, fb, fd, fe)
        // Here we do not need fb as it is only used to mask some values before sending
        // them to other participants then adding them all together to generate w.
        // This sum would annihilate all the fb shares which make them useless in our case.
        let fk = Polynomial::generate_polynomial(None, max_malicious, rng).unwrap();
        let fa = Polynomial::generate_polynomial(None, max_malicious, rng).unwrap();
        let degree = 2usize.checked_mul(max_malicious).unwrap();
        let fd = Polynomial::generate_polynomial(Some(Secp256K1ScalarField::zero()), degree, rng)
            .unwrap();
        let fe = Polynomial::generate_polynomial(Some(Secp256K1ScalarField::zero()), degree, rng)
            .unwrap();

        // computing k, R, Rx
        let k = fk.eval_at_zero().unwrap().0;
        let big_r = ProjectivePoint::GENERATOR * k;

        // compute the master scalar w = a * k
        let w = fa.eval_at_zero().unwrap().0 * k;
        let w_invert = w.invert().unwrap();

        (w_invert, fa, fd, fe, big_r)
    }

    fn simulate_sign_inputs(
        participants: &[Participant],
        max_malicious: usize,
        rng: &mut impl CryptoRngCore,
    ) -> (ProjectivePoint, Vec<(Participant, PresignOutput)>) {
        let fx = Polynomial::generate_polynomial(None, max_malicious, rng).unwrap();
        let x = fx.eval_at_zero().unwrap().0;
        let public_key = ProjectivePoint::GENERATOR * x;
        let (w_invert, fa, fd, fe, big_r) = simulate_presignature(max_malicious, rng);
        let presignatures = participants
            .iter()
            .map(|p| {
                let c_i = w_invert * fa.eval_at_participant(*p).unwrap().0;
                let alpha = c_i + fd.eval_at_participant(*p).unwrap().0;
                let beta = c_i * fx.eval_at_participant(*p).unwrap().0;
                let e = fe.eval_at_participant(*p).unwrap().0;
                (
                    *p,
                    PresignOutput {
                        big_r: big_r.to_affine(),
                        alpha,
                        beta,
                        e,
                        c: c_i,
                    },
                )
            })
            .collect();
        (public_key, presignatures)
    }

    #[test]
    fn test_sign_given_presignature_without_rerandomization() {
        let mut rng = MockCryptoRng::seed_from_u64(42);
        let max_malicious = 2;
        let msg = b"Hello? Is it me you're looking for?";

        let participants = generate_participants(5);
        let (public_key, participants_presign) =
            simulate_sign_inputs(&participants, max_malicious, &mut rng);

        let (_, sig) = run_sign_without_rerandomization(
            &participants_presign,
            max_malicious.into(),
            public_key,
            msg,
            &mut rng,
        )
        .unwrap();
        let sig = ecdsa::Signature::from_scalars(x_coordinate(&sig.big_r), sig.s).unwrap();

        // verify the correctness of the generated signature
        VerifyingKey::from(&PublicKey::from_affine(public_key.to_affine()).unwrap())
            .verify(&msg[..], &sig)
            .unwrap();
    }

    #[test]
    fn test_sign_given_presignature_with_rerandomization() {
        let mut rng = MockCryptoRng::seed_from_u64(42);
        let max_malicious = 2;
        let msg = b"Hello? Is it me you're looking for?";

        let participants = generate_participants(5);
        let (public_key, participants_presign) =
            simulate_sign_inputs(&participants, max_malicious, &mut rng);
        let public_key_vk = frost_core::VerifyingKey::new(public_key);

        let (tweak, _, sig) = run_sign_with_rerandomization(
            &participants_presign,
            max_malicious,
            public_key,
            msg,
            &mut rng,
        )
        .unwrap();
        let signature = ecdsa::Signature::from_scalars(x_coordinate(&sig.big_r), sig.s).unwrap();
        let public_key = tweak.derive_verifying_key(&public_key_vk).to_element();

        // verify the correctness of the generated signature
        VerifyingKey::from(&PublicKey::from_affine(public_key.to_affine()).unwrap())
            .verify(&msg[..], &signature)
            .unwrap();

        insta::assert_json_snapshot!(signature);
    }

    #[test]
    fn test_sign_fails_if_s_is_zero() {
        let mut rng = MockCryptoRng::seed_from_u64(42);
        let participants = generate_participants(3);
        let max_malicious = 1;

        // presignatures with s_me = 0 for each participant
        let presignatures = participants
            .iter()
            .map(|p| {
                (
                    *p,
                    PresignOutput {
                        big_r: ProjectivePoint::IDENTITY.to_affine(),
                        alpha: Secp256K1ScalarField::zero(),
                        beta: Secp256K1ScalarField::zero(),
                        c: Secp256K1ScalarField::zero(),
                        e: Secp256K1ScalarField::zero(),
                    },
                )
            })
            .collect::<Vec<_>>();

        let public_key = ProjectivePoint::IDENTITY;
        let msg = [0u8; 32]; // arbitrary zero message

        let result = crate::ecdsa::robust_ecdsa::test::run_sign_without_rerandomization(
            &presignatures,
            max_malicious.into(),
            public_key,
            &msg,
            &mut rng,
        );

        match result {
            Ok(_) => panic!("expected failure, got success"),
            Err(err) => {
                let text = err.to_string();
                assert!(
                    text.contains("signature part s cannot be zero"),
                    "unexpected error type: {text}"
                );
            }
        }
    }

    #[test]
    fn test_sign_fails_if_max_malicious_too_high() {
        let mut rng = MockCryptoRng::seed_from_u64(42);
        let participants = generate_participants(2);
        let max_malicious = 1;

        // arbitrary values
        let presignatures = participants
            .iter()
            .map(|p| {
                (
                    *p,
                    PresignOutput {
                        big_r: ProjectivePoint::IDENTITY.to_affine(),
                        alpha: Secp256K1ScalarField::zero(),
                        beta: Secp256K1ScalarField::zero(),
                        c: Secp256K1ScalarField::zero(),
                        e: Secp256K1ScalarField::zero(),
                    },
                )
            })
            .collect::<Vec<_>>();
        let public_key = ProjectivePoint::IDENTITY;
        let msg = [0u8; 32];

        let result = crate::ecdsa::robust_ecdsa::test::run_sign_without_rerandomization(
            &presignatures,
            max_malicious.into(),
            public_key,
            &msg,
            &mut rng,
        );

        match result {
            Ok(_) => panic!("expected failure, got success"),
            Err(err) => {
                let text = err.to_string();
                assert!(
                    text.contains("bad parameters: 2*max_malicious+1 must be less than or equals to participant count"),
                    "unexpected error type: {text}"
                );
            }
        }
    }

    #[rstest]
    #[case(1)]
    #[case(2)]
    #[case(3)]
    fn test_sign_buffer_entries(#[case] max_malicious: usize) {
        // Given
        let mut rng = MockCryptoRng::seed_from_u64(42);
        let num_participants = 2 * max_malicious + 1;
        let participants = generate_participants(num_participants);
        let (public_key, presignatures) =
            simulate_sign_inputs(&participants, max_malicious, &mut rng);
        let coordinator = participants[0];
        let msg_scalar = crate::crypto::hash::test::scalar_hash_secp256k1(b"test msg");

        // When + Then
        assert_buffer_capacity(
            &participants,
            &mut rng,
            |comms, p_list, p, _rng_p| {
                let presignature = &presignatures.iter().find(|(pp, _)| *pp == p).unwrap().1;
                let rerandomized =
                    crate::ecdsa::robust_ecdsa::RerandomizedPresignOutput::new_without_rerandomization(
                        presignature,
                    );
                fut_wrapper(
                    comms.shared_channel(),
                    p_list,
                    coordinator,
                    p,
                    public_key.to_affine(),
                    rerandomized,
                    msg_scalar,
                )
            },
            expected_buffer_by_role(
                coordinator,
                ROBUST_ECDSA_SIGN_MAX_INCOMING_COORDINATOR_ENTRIES,
                ROBUST_ECDSA_SIGN_MAX_INCOMING_PARTICIPANT_ENTRIES,
            ),
        );
    }
}