use crate::member::Member;
use crate::signature::Signature;
use curve25519_dalek::ristretto::CompressedRistretto;
use curve25519_dalek::scalar::Scalar;

use crate::transcript::TranscriptProtocol;
use merlin::Transcript;

// This module will pull together all of the necessary things
// Setting up everything we need
#[derive(Debug)]
pub enum Error {
    // This error occurs when the sign method is called
    // without a signer being in the ring
    NoSigner,
    // This error occurs if there are less than 2 members in the ring
    NotEnoughMembers,
    // This error occurs if all members do not have the same number of keys
    NumberOfKeysMismatch,
    // This error occurs if there is more than one signer in the ring
    MoreThanOneSigner,
    // This error occurs if a member in the ring has duplicate keys
    DuplicateKeysExist,
    // This error occurs when an underlying module produces an error
    UnderlyingErr(String),
}

impl From<crate::member::Error> for crate::clsag::Error {
    fn from(e: crate::member::Error) -> crate::clsag::Error {
        match e {
            crate::member::Error::NotASigner => Error::UnderlyingErr(String::from(
                "Tried to use a method specific to a signer in the clsag module",
            )),
            crate::member::Error::NotADecoy => Error::UnderlyingErr(String::from(
                "Tried to use a method specific to a decoy in the clsag module",
            )),
        }
    }
}
// This struct is used to construct the clsag signature
pub struct Clsag {
    members: Vec<Member>,
}

impl Clsag {
    // Creates a new clsag component with a configured basepoint
    pub fn new() -> Self {
        Clsag {
            members: Vec::new(),
        }
    }
    // Adds a member to the clsag component
    // Use this method to add decoys and signers to the struct
    pub fn add_member(&mut self, member: Member) {
        self.members.push(member);
    }
    // Returns public keys from all known members as vector of bytes
    pub fn public_keys_bytes(&self) -> Vec<u8> {
        self.members
            .iter()
            .map(|member| member.public_set.to_bytes())
            .flatten()
            .collect()
    }
    // Returns public keys from all known members
    pub fn public_keys(&self) -> Vec<Vec<CompressedRistretto>> {
        self.members
            .iter()
            .map(|member| member.public_set.to_keys())
            .collect()
    }
    // sign produces a clsag signature
    pub fn sign(&self) -> Result<Signature, Error> {
        self.check_format()?;

        let num_members = self.members.len();
        let mut all_challenges: Vec<Scalar> = Vec::with_capacity(num_members);

        // Fetch signer of the ring
        let signer_index = self.find_signer()?;
        let signer = &self.members[signer_index];

        // Fetch public key matrix
        let pubkey_matrix = self.public_keys_bytes();

        // Compute key images for signer
        let key_images = signer.compute_key_images()?;

        // Calculate aggrgation co-efficients
        let aggregation_cooeff = calc_aggregation_coefficients(&pubkey_matrix, &key_images);

        // Calculate first challenge
        let mut challenge = signer.compute_challenge_commitment(&pubkey_matrix)?;
        all_challenges.push(challenge);

        // seed challenge into for loop starting from member after signer
        for decoy in self
            .members
            .iter()
            .cycle()
            .skip(signer_index + 1)
            .take(num_members - 1)
        {
            challenge = decoy.compute_decoy_challenge(
                &challenge,
                &key_images,
                &aggregation_cooeff,
                &pubkey_matrix,
            )?;
            all_challenges.push(challenge);
        }

        // The last challenge variable should be the one generated by the member before the signer,
        // which will be for the signer. The signer will use this to generate his response values
        // and close the ring
        let signers_response = signer.compute_signer_response(challenge, &aggregation_cooeff)?;

        // Collect all responses
        let mut all_responses: Vec<Scalar> = Vec::with_capacity(num_members);
        for member in &self.members {
            match member.is_signer() {
                true => {
                    all_responses.push(signers_response);
                }
                false => {
                    let mem_response =
                        member
                            .response
                            .clone()
                            .ok_or(Error::UnderlyingErr(String::from(
                                "member does not have a response value",
                            )))?;
                    all_responses.push(mem_response);
                }
            }
        }

        // Collect first members challenge
        // The last element in the vector of challenges will be the signers challenge
        // Since we also know the index `n` of the signer. We need to walk back `n` times
        // This is equivalent to reversing the vector and getting the `n`th element
        // to fetch the first challenge; the challenge corresponding to the first member
        all_challenges.reverse();
        let first_challenge = all_challenges[signer_index];

        Ok(Signature {
            challenge: first_challenge,
            responses: all_responses,
            key_images: key_images,
        })
    }
    // Returns the position of the signer in the ring
    // If this call is completed after check_format, it should not fail
    // as check_format ensures there is one signer. This method has been
    // added for a cleaner API. The alternative would be for the method which checks
    // that clsag is formatted correctly, to also return the signer.
    pub fn find_signer(&self) -> Result<usize, Error> {
        let signer_index = self
            .members
            .iter()
            .position(|member| member.is_signer())
            .ok_or(Error::NoSigner)?;

        Ok(signer_index)
    }
    // Returns the number of signers in the ring
    pub fn num_signers(&self) -> usize {
        let signers: Vec<&Member> = self
            .members
            .iter()
            .filter(|member| member.is_signer())
            .collect();
        signers.len()
    }
    // Checks that the clsag is correctly constructed
    fn check_format(&self) -> Result<(), Error> {
        // Check that we have more than one member
        if self.members.len() < 2 {
            return Err(Error::NotEnoughMembers);
        }

        // Check there is only one signer in the ring
        let num_signers = self.num_signers();
        match num_signers {
            0 => return Err(Error::NoSigner),
            1 => (),
            _ => return Err(Error::MoreThanOneSigner),
        };

        // Check that each member has the same number of keys
        let first_member_num_keys = self.members[0].num_keys();
        let all_same_num_keys = self
            .members
            .iter()
            .all(|member| member.num_keys() == first_member_num_keys);
        if !all_same_num_keys {
            return Err(Error::NumberOfKeysMismatch);
        }

        // Check that each member has no duplicates
        let no_duplicates_exists = self
            .members
            .iter()
            .all(|member| !member.public_set.duplicates_exist());
        if !no_duplicates_exists {
            return Err(Error::DuplicateKeysExist);
        }
        Ok(())
    }
}
// Calculates the aggregation co-efficients mu_x and [mu_j]
pub fn calc_aggregation_coefficients(
    pubkey_matrix: &[u8],
    key_images: &[CompressedRistretto],
) -> Vec<Scalar> {
    // precompute (pubkey_matrix || keyimages)
    let key_images_bytes: Vec<u8> = key_images
        .iter()
        .map(|key_image| key_image.to_bytes().to_vec())
        .flatten()
        .collect();

    let num_keys_per_user = key_images.len();

    let mut agg_coef = Vec::with_capacity(num_keys_per_user);

    for i in 0..num_keys_per_user {
        let mut transcript = Transcript::new(b"clsag");

        transcript.append_u64(b"", i as u64);
        transcript.append_message(b"", pubkey_matrix);
        transcript.append_message(b"", &key_images_bytes);

        agg_coef.push(transcript.challenge_scalar(b""));
    }

    agg_coef
}

#[cfg(test)]
mod test {
    extern crate test;

    use super::*;
    use crate::tests_helper::*;
    use test::Bencher;

    #[test]
    fn test_check_format() {
        let num_decoys = 10;
        let num_keys = 3;
        let mut clsag = generate_clsag_with(num_decoys, num_keys);

        // No signer in the ring
        match clsag.sign() {
            Ok(_) => panic!("expected an error as there is no signer in the ring"),
            Err(Error::NoSigner) => {}
            Err(_) => panic!("got an error, however we expected no signer error"),
        }

        // Add a signer
        clsag.add_member(generate_signer(num_keys));
        // Another one
        clsag.add_member(generate_signer(num_keys));

        // More than one signer in the ring
        match clsag.sign() {
            Ok(_) => panic!("expected an error as there are too many signers in the ring"),
            Err(Error::MoreThanOneSigner) => {}
            Err(_) => panic!("got an error, however we expected a more than one signer error"),
        }

        clsag = generate_clsag_with(num_decoys, num_keys);
        // Add different number of keys
        clsag.add_member(generate_decoy(num_keys + 1));

        // Add correct signer
        clsag.add_member(generate_signer(num_keys));

        // One member has a different number of keys
        match clsag.sign() {
            Ok(_) => {
                panic!("expected an error as one member has more keys than another in the ring")
            }
            Err(Error::NumberOfKeysMismatch) => {}
            Err(_) => panic!("got an error, however we expected a `number of keys mismatch` error"),
        };

        clsag = generate_clsag_with(num_decoys, num_keys);
        // Add correct signer
        clsag.add_member(generate_signer(num_keys));

        // Set the first key in members key set to the value of the last key
        let first_member = &mut clsag.members[0];
        let first_member_last_element = &mut first_member.public_set.0.last().unwrap();
        first_member.public_set.0[0] = first_member_last_element.clone();

        match clsag.sign() {
            Ok(_) => panic!("expected an error as one member has a duplicate key"),
            Err(Error::DuplicateKeysExist) => {}
            Err(_) => panic!("got an error, however we expected a `duplicate keys` error"),
        };
    }

    #[test]
    fn test_sign_no_error() {
        let num_decoys = 10;
        let num_keys = 3;
        let mut clsag = generate_clsag_with(num_decoys, num_keys);

        // Add a signer
        clsag.add_member(generate_signer(num_keys));

        // Should produce no error
        let signature = clsag.sign().unwrap();

        // number of key images should equal number of keys
        assert_eq!(num_keys, signature.key_images.len());

        // number of responses should equal number of key images * number of members
        // number of members = number of decoys + signer
        let num_members = num_decoys + 1;
        assert_eq!(num_members, signature.responses.len());
    }

    #[bench]
    fn bench_sign(b: &mut Bencher) {
        // One time setup code here
        let num_keys = 2;
        let num_decoys = 11;

        let mut clsag = generate_clsag_with(num_decoys, num_keys);
        clsag.add_member(generate_signer(num_keys));

        b.iter(|| clsag.sign());
    }
}