#![allow(non_snake_case)]

//! Simple ed25519
//!
//! See https://tools.ietf.org/html/rfc8032
//! This is an implementation of the Musig2 protocol as shown in https://eprint.iacr.org/2020/1261.pdf with the addition named Musig2* suggested in Section B of the paper.
//! We implement the v = 2 (NUMBER_OF_NONCES) version, meaning there are 2 nonces generated by each party.

use super::{ExpandedKeyPair, Signature};
use curv::arithmetic::Converter;
use curv::elliptic::curves::{Ed25519, Point, Scalar};
use curv::BigInt;
use protocols::Rng;
use sha2::{digest::Digest, Sha512};

pub const NUMBER_OF_NONCES: usize = 2;

#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
pub struct PublicKeyAgg {
    pub agg_public_key: Point<Ed25519>,
    pub musig_coefficient: Scalar<Ed25519>,
}

impl PublicKeyAgg {
    #[allow(clippy::unnecessary_sort_by)]
    pub fn key_aggregation_n(
        mut public_keys: Vec<Point<Ed25519>>,
        my_public_key: &Point<Ed25519>,
    ) -> Option<PublicKeyAgg> {
        // When there are at least 2 distinct public keys, it is secure to set the musig coefficient
        // of one them to 1 - saving a scalar multiplication operation - proof in Section B of the Musig2 paper linked above.
        // We therefore find the second public key (by lexicographic order) and later set its musig coefficient to 1.
        public_keys.sort_by(|left, right| left.to_bytes(false).cmp(&right.to_bytes(false)));
        let mut second_public_key = &public_keys[0];
        for public_key in &public_keys[1..] {
            if *public_key.to_bytes(false) > *public_keys[0].to_bytes(false) {
                second_public_key = public_key;
                break;
            }
        }
        let mut found_my_pub_key = false;
        let (sum, my_coeff) = public_keys.iter().fold(
            (Point::<Ed25519>::zero(), Scalar::<Ed25519>::from(1)),
            |(mut agg_pub_key, mut musig_coeff), public_key| {
                let mut musig_coefficient: Scalar<Ed25519> = Scalar::from(1);
                if public_key != second_public_key {
                    let mut hasher = Sha512::new().chain(&[1]).chain(&*public_key.to_bytes(true));
                    for pk in &public_keys {
                        hasher.update(&*pk.to_bytes(true));
                    }
                    let mut hash_result = hasher.finalize();
                    // reverse because BigInt uses BigEndian.
                    hash_result.reverse();
                    // This will reduce it mod the group order.
                    musig_coefficient = Scalar::from_bigint(&BigInt::from_bytes(&hash_result));
                }

                let party_i_public_key = public_key * &musig_coefficient;
                if public_key == my_public_key {
                    musig_coeff = musig_coefficient;
                    found_my_pub_key = true;
                }
                agg_pub_key = &agg_pub_key + party_i_public_key;
                (agg_pub_key, musig_coeff)
            },
        );
        if found_my_pub_key {
            Some(PublicKeyAgg {
                agg_public_key: sum,
                musig_coefficient: my_coeff,
            })
        } else {
            None
        }
    }
}
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
pub struct PrivatePartialNonces {
    pub r: [Scalar<Ed25519>; NUMBER_OF_NONCES],
}

#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
pub struct PublicPartialNonces {
    pub R: [Point<Ed25519>; NUMBER_OF_NONCES],
}

#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
pub struct PartialSignature {
    pub R: Point<Ed25519>,
    pub my_partial_s: Scalar<Ed25519>,
}

pub fn generate_partial_nonces(
    keys: &ExpandedKeyPair,
    message: Option<&[u8]>,
) -> (PrivatePartialNonces, PublicPartialNonces) {
    let mut rng = rand::thread_rng();
    generate_partial_nonces_internal(keys, message, &mut rng)
}

fn generate_partial_nonces_internal(
    keys: &ExpandedKeyPair,
    message: Option<&[u8]>,
    rng: &mut impl Rng,
) -> (PrivatePartialNonces, PublicPartialNonces) {
    // here we deviate from the spec, by introducing  non-deterministic element (random number)
    // to the nonce, this is important for MPC implementations
    let r: [Scalar<Ed25519>; NUMBER_OF_NONCES] = [(); NUMBER_OF_NONCES].map(|_| {
        let mut hash_result = Sha512::new()
            .chain(&[2])
            .chain(&*keys.expanded_private_key.prefix.to_bytes())
            .chain(message.unwrap_or(&[]))
            .chain(rng.gen::<[u8; 32]>())
            .finalize();
        // reverse because BigInt uses big-endian
        hash_result.reverse();
        // reduce modulu the group order
        Scalar::from_bigint(&BigInt::from_bytes(&hash_result))
    });
    let R: [Point<Ed25519>; NUMBER_OF_NONCES] = r.clone().map(|scalar| Point::generator() * scalar);
    (PrivatePartialNonces { r }, PublicPartialNonces { R })
}

pub fn partial_sign(
    nonces_from_other_parties: &[[Point<Ed25519>; NUMBER_OF_NONCES]],
    my_private_partial_nonces: PrivatePartialNonces,
    my_public_partial_nonces: PublicPartialNonces,
    agg_public_key: &PublicKeyAgg,
    my_keypair: &ExpandedKeyPair,
    message: &[u8],
) -> PartialSignature {
    // Sum up the partial nonces from all parties index-wise, meaning,  R[i]
    // is the sum of partial_nonces[i] from all parties
    let R: [Point<Ed25519>; NUMBER_OF_NONCES] = nonces_from_other_parties.iter().fold(
        my_public_partial_nonces.R,
        |mut sum_partial_nonces, partial_nonce_array| {
            for (accum_nonce, nonce) in sum_partial_nonces.iter_mut().zip(partial_nonce_array) {
                *accum_nonce = &*accum_nonce + nonce;
            }
            sum_partial_nonces
        },
    );

    // Compute b as hash of nonces
    let mut hasher = Sha512::new()
        .chain(&[3])
        .chain(&*agg_public_key.agg_public_key.to_bytes(false));
    for nonce in &R {
        hasher.update(&*nonce.to_bytes(false));
    }
    hasher.update(message);
    let mut hash_result = hasher.finalize();
    // Reverse because BigInt uses big-endian
    hash_result.reverse();
    // Reduce modulu the group order
    let b: Scalar<Ed25519> = Scalar::from_bigint(&BigInt::from_bytes(&hash_result));
    // Compute effective nonce
    // The idea is to compute R and r s.t. R = R_0 + b•R_1 + ... + b^(v-1)•R_v and r = r_0 + b•r_1 + ... + b^(v-1)•r_v
    let (effective_R, effective_r, _) = R[1..]
        .iter()
        .zip(my_private_partial_nonces.r[1..].iter())
        .fold(
            (
                R[0].clone(),
                my_private_partial_nonces.r[0].clone(),
                b.clone(),
            ),
            |(eff_R, eff_r, b_exp), (nonce_R_i, nonce_r_i)| {
                (
                    eff_R + &b_exp * nonce_R_i,
                    &eff_r + &b_exp * nonce_r_i,
                    b_exp * &b,
                )
            },
        );
    // Compute Fiat-Shamir challenge of signature
    let sig_challenge = Signature::k(&effective_R, &agg_public_key.agg_public_key, message);

    // Computes the partial signature
    let partial_signature: Scalar<Ed25519> = sig_challenge
        * &agg_public_key.musig_coefficient
        * &my_keypair.expanded_private_key.private_key
        + effective_r;
    PartialSignature {
        R: effective_R,
        my_partial_s: partial_signature,
    }
}

pub fn aggregate_partial_signatures(
    my_partial_sig: &PartialSignature,
    partial_sigs_from_other_parties: &[Scalar<Ed25519>],
) -> Signature {
    let aggregate_signature = partial_sigs_from_other_parties
        .iter()
        .sum::<Scalar<Ed25519>>()
        + &my_partial_sig.my_partial_s;

    Signature {
        R: my_partial_sig.R.clone(),
        s: aggregate_signature,
    }
}

mod test;
