use crate::prelude::*;

use curve25519_dalek::constants;
use curve25519_dalek::ristretto::RistrettoPoint;
use curve25519_dalek::scalar::Scalar;
use curve25519_dalek::traits::MultiscalarMul;
use digest::Digest;
use digest::generic_array::typenum::U64;
use rand_core::{CryptoRng, RngCore};

use crate::traits::{KeyImageGen, Link, Sign, Verify};

/// Dual Linkable Spontaneous Anonymous Group Signature for Ad Hoc Groups
///
/// [DLSAG: Non-Interactive Refund Transactions For Interoperable Payment Channels in Monero](https://eprint.iacr.org/2019/595.pdf)
///
/// This is a a novel linkable ring signature scheme that enables for the
/// first time refund transactions natively in Monero
///
/// Read the paper on how to use it to implement payment channels.
///
/// Please read tests at the bottom of the source code for this module for
/// examples on how to use it
#[derive(Clone)]
pub struct DLSAG {
    pub challenge: Scalar,
    pub responses: Vec<Scalar>,
    pub ring: Vec<(RistrettoPoint, RistrettoPoint, Scalar)>,
    pub key_image: RistrettoPoint,
    pub b: bool,
}

impl KeyImageGen<(Scalar, RistrettoPoint, Scalar), RistrettoPoint> for DLSAG {
    /// Some signature schemes require the key images to be signed as well.
    /// Use this method to generate them
    fn generate_key_image<Hash: Digest<OutputSize=U64> + Clone + Default>(
        k: (Scalar, RistrettoPoint, Scalar),
    ) -> RistrettoPoint {
        let k_point: (RistrettoPoint, RistrettoPoint, Scalar) =
            (k.0 * constants::RISTRETTO_BASEPOINT_POINT, k.1, k.2);

        let key_image: RistrettoPoint = k.2
            * k.0
            * RistrettoPoint::from_hash(Hash::default().chain(k_point.1.compress().as_bytes()));

        return key_image;
    }
}

impl KeyImageGen<(RistrettoPoint, Scalar, Scalar), RistrettoPoint> for DLSAG {
    /// Some signature schemes require the key images to be signed as well.
    /// Use this method to generate them
    fn generate_key_image<Hash: Digest<OutputSize=U64> + Clone + Default>(
        k: (RistrettoPoint, Scalar, Scalar),
    ) -> RistrettoPoint {
        let k_point: (RistrettoPoint, RistrettoPoint, Scalar) =
            (k.0, k.1 * constants::RISTRETTO_BASEPOINT_POINT, k.2);

        let key_image: RistrettoPoint = k.2
            * k.1
            * RistrettoPoint::from_hash(Hash::default().chain(k_point.0.compress().as_bytes()));

        return key_image;
    }
}

impl Sign<(Scalar, RistrettoPoint, Scalar), Vec<(RistrettoPoint, RistrettoPoint, Scalar)>>
for DLSAG
{
    /// To sign you need `k` your private key, and `ring` which is the public keys of everyone
    /// except you. You are signing the `message`
    ///
    /// The private key `k` in this case is your private key, the public key of the other end of
    /// the channel and a random bitstring generated by hashing-to-scalar: the transaction ID, and
    /// output index.
    ///
    /// The ring contains public key pairs from the blockchain together with their random
    /// bitstrings as mentioned above.
    ///
    /// This implementation of `sign(...)` is for the user who has the private key for the left
    /// side of the channel
    fn sign<
        Hash: Digest<OutputSize=U64> + Clone + Default,
        CSPRNG: CryptoRng + RngCore + Default,
    >(
        k: (Scalar, RistrettoPoint, Scalar),
        mut ring: Vec<(RistrettoPoint, RistrettoPoint, Scalar)>,
        secret_index: usize,
        message: &Vec<u8>,
    ) -> DLSAG {
        let mut csprng = CSPRNG::default();

        // Provers public key
        let k_point: (RistrettoPoint, RistrettoPoint, Scalar) =
            (k.0 * constants::RISTRETTO_BASEPOINT_POINT, k.1, k.2);

        let key_image: RistrettoPoint = DLSAG::generate_key_image::<Hash>(k);

        // Ring size (at least 4 but maximum 32)
        let n = ring.len() + 1;

        ring.insert(secret_index, k_point);

        let a: Scalar = Scalar::random(&mut csprng);

        let mut rs: Vec<Scalar> = (0..n).map(|_| Scalar::random(&mut csprng)).collect();

        let mut cs: Vec<Scalar> = (0..n).map(|_| Scalar::zero()).collect();

        // Hash of message is shared by all challenges H_n(m, ....)
        let mut message_hash = Hash::default();

        message_hash.update(message);

        let mut hashes: Vec<Hash> = (0..n).map(|_| message_hash.clone()).collect();

        hashes[(secret_index + 1) % n].update(
            (a * constants::RISTRETTO_BASEPOINT_POINT)
                .compress()
                .as_bytes(),
        );
        hashes[(secret_index + 1) % n].update(
            (a * ring[secret_index].2
                * RistrettoPoint::from_hash(
                Hash::default().chain(k_point.1.compress().as_bytes()),
            ))
                .compress()
                .as_bytes(),
        );
        cs[(secret_index + 1) % n] = Scalar::from_hash(hashes[(secret_index + 1) % n].clone());

        let mut i = (secret_index + 1) % n;

        loop {
            hashes[(i + 1) % n].update(
                RistrettoPoint::multiscalar_mul(
                    &[rs[i % n], cs[i % n]],
                    &[constants::RISTRETTO_BASEPOINT_POINT, ring[i % n].0],
                )
                    .compress()
                    .as_bytes(),
            );
            hashes[(i + 1) % n].update(
                RistrettoPoint::multiscalar_mul(
                    &[rs[i % n], cs[i % n]],
                    &[
                        ring[i % n].2 * RistrettoPoint::from_hash(
                            Hash::default().chain(
                                ring[i % n].1.compress().as_bytes()
                            )
                        ),
                        key_image
                    ],
                )
                    .compress()
                    .as_bytes(),
            );
            cs[(i + 1) % n] = Scalar::from_hash(hashes[(i + 1) % n].clone());

            if secret_index >= 1 && i % n == (secret_index - 1) % n {
                break;
            } else if secret_index == 0 && i % n == n - 1 {
                break;
            } else {
                i = (i + 1) % n;
            }
        }

        rs[secret_index] = a - (cs[secret_index] * k.0);

        return DLSAG {
            challenge: cs[0],
            responses: rs,
            ring: ring,
            key_image: key_image,
            b: false,
        };
    }
}

impl Sign<(RistrettoPoint, Scalar, Scalar), Vec<(RistrettoPoint, RistrettoPoint, Scalar)>>
for DLSAG
{
    /// To sign you need `k` your private key, and `ring` which is the public keys of everyone
    /// except you. You are signing the `message`
    ///
    /// The private key `k` in this case is your private key, the public key of the other end of
    /// the channel and a random bitstring generated by hashing-to-scalar: the transaction ID, and
    /// output index.
    ///
    /// The ring contains public key pairs from the blockchain together with their random
    /// bitstrings as mentioned above.
    ///
    /// This implementation of `sign(...)` is for the user who has the private key for the right
    /// side of the channel
    fn sign<
        Hash: Digest<OutputSize=U64> + Clone + Default,
        CSPRNG: CryptoRng + RngCore + Default,
    >(
        k: (RistrettoPoint, Scalar, Scalar),
        mut ring: Vec<(RistrettoPoint, RistrettoPoint, Scalar)>,
        secret_index: usize,
        message: &Vec<u8>,
    ) -> DLSAG {
        let mut csprng = CSPRNG::default();

        // Provers public key
        let k_point: (RistrettoPoint, RistrettoPoint, Scalar) =
            (k.0, k.1 * constants::RISTRETTO_BASEPOINT_POINT, k.2);

        let key_image: RistrettoPoint = DLSAG::generate_key_image::<Hash>(k);

        // Ring size (at least 4 but maximum 32)
        let n = ring.len() + 1;

        ring.insert(secret_index, k_point);

        let a: Scalar = Scalar::random(&mut csprng);

        let mut rs: Vec<Scalar> = (0..n).map(|_| Scalar::random(&mut csprng)).collect();

        let mut cs: Vec<Scalar> = (0..n).map(|_| Scalar::zero()).collect();

        // Hash of message is shared by all challenges H_n(m, ....)
        let mut message_hash = Hash::default();

        message_hash.update(message);

        let mut hashes: Vec<Hash> = (0..n).map(|_| message_hash.clone()).collect();

        hashes[(secret_index + 1) % n].update(
            (a * constants::RISTRETTO_BASEPOINT_POINT)
                .compress()
                .as_bytes(),
        );
        hashes[(secret_index + 1) % n].update(
            (a * ring[secret_index].2
                * RistrettoPoint::from_hash(
                Hash::default().chain(k_point.0.compress().as_bytes()),
            ))
                .compress()
                .as_bytes(),
        );
        cs[(secret_index + 1) % n] = Scalar::from_hash(hashes[(secret_index + 1) % n].clone());

        let mut i = (secret_index + 1) % n;

        loop {
            hashes[(i + 1) % n].update(
                RistrettoPoint::multiscalar_mul(
                    &[rs[i % n], cs[i % n]],
                    &[
                        constants::RISTRETTO_BASEPOINT_POINT,
                        ring[i % n].1
                    ],
                )
                    .compress()
                    .as_bytes(),
            );
            hashes[(i + 1) % n].update(
                RistrettoPoint::multiscalar_mul(
                    &[rs[i % n], cs[i % n]],
                    &[
                        ring[i % n].2 * RistrettoPoint::from_hash(
                            Hash::default().chain(
                                ring[i % n].0.compress().as_bytes()
                            ),
                        ),
                        key_image
                    ],
                )
                    .compress()
                    .as_bytes(),
            );
            cs[(i + 1) % n] = Scalar::from_hash(hashes[(i + 1) % n].clone());

            if secret_index >= 1 && i % n == (secret_index - 1) % n {
                break;
            } else if secret_index == 0 && i % n == n - 1 {
                break;
            } else {
                i = (i + 1) % n;
            }
        }

        rs[secret_index] = a - (cs[secret_index] * k.1);

        return DLSAG {
            challenge: cs[0],
            responses: rs,
            ring: ring,
            key_image: key_image,
            b: true,
        };
    }
}

impl Verify for DLSAG {
    /// To verify a `signature` you need the `message` too
    fn verify<Hash: Digest<OutputSize=U64> + Clone + Default>(
        signature: DLSAG,
        message: &Vec<u8>,
    ) -> bool {
        let mut reconstructed_c: Scalar = signature.challenge;
        let n = signature.ring.len();
        for j in 0..n {
            let mut h: Hash = Hash::default();
            h.update(message);
            if signature.b {
                h.update(
                    RistrettoPoint::multiscalar_mul(
                        &[signature.responses[j], reconstructed_c],
                        &[
                            constants::RISTRETTO_BASEPOINT_POINT,
                            signature.ring[j].1
                        ],
                    )
                        .compress()
                        .as_bytes(),
                );

                h.update(
                    RistrettoPoint::multiscalar_mul(
                        &[signature.responses[j], reconstructed_c],
                        &[
                            signature.ring[j].2 * RistrettoPoint::from_hash(
                            Hash::default().chain(
                                    signature.ring[j].0.compress().as_bytes()
                                ),
                            ),
                            signature.key_image
                        ]
                    )
                        .compress()
                        .as_bytes(),
                );
            } else {
                h.update(
                    RistrettoPoint::multiscalar_mul(
                        &[signature.responses[j], reconstructed_c],
                        &[
                            constants::RISTRETTO_BASEPOINT_POINT,
                            signature.ring[j].0
                        ]
                    )
                        .compress()
                        .as_bytes(),
                );

                h.update(
                    RistrettoPoint::multiscalar_mul(
                        &[signature.responses[j], reconstructed_c],
                        &[
                            signature.ring[j].2 * RistrettoPoint::from_hash(
                                Hash::default().chain(
                                    signature.ring[j].1.compress().as_bytes()
                                )
                            ),
                            signature.key_image
                        ])
                        .compress()
                        .as_bytes(),
                );
            }
            reconstructed_c = Scalar::from_hash(h);
        }

        return signature.challenge == reconstructed_c;
    }
}

impl Link for DLSAG {
    /// This is for linking two signatures and checking if they are signed by the same person
    fn link(signature_1: DLSAG, signature_2: DLSAG) -> bool {
        return signature_1.key_image == signature_2.key_image;
    }
}

#[cfg(test)]
#[cfg(feature = "std")]
mod test {
    extern crate blake2;
    extern crate rand;
    extern crate sha2;
    extern crate sha3;

    use blake2::Blake2b;
    use curve25519_dalek::ristretto::RistrettoPoint;
    use curve25519_dalek::scalar::Scalar;
    use rand::rngs::OsRng;
    use sha2::Sha512;
    use sha3::Keccak512;

    use super::*;

    #[test]
    fn dlsag() {
        let mut csprng = OsRng::default();
        let k: (Scalar, RistrettoPoint, Scalar) = (
            Scalar::random(&mut csprng),         // The prover's private key
            RistrettoPoint::random(&mut csprng), // The public key of the other end of the channel
            // According to the paper this should be a random
            // bitstring generated by hashing transaction ID, and output index.
            // It is simulated here using a random bitstring
            Scalar::random(&mut csprng),
        );

        let other_k: (RistrettoPoint, Scalar, Scalar) = (k.1, k.0, k.2);
        let secret_index = 1;
        let n = 2;
        // Simulate randomly chosen Public keys (Prover will insert her public key here later)
        let ring: Vec<(RistrettoPoint, RistrettoPoint, Scalar)> = (0..(n - 1)) // Prover is going to add her key into this mix
            .map(|_| {
                (
                    RistrettoPoint::random(&mut csprng),
                    RistrettoPoint::random(&mut csprng),
                    Scalar::random(&mut csprng),
                )
            })
            .collect();
        let message: Vec<u8> = b"This is the message".iter().cloned().collect();

        {
            let signature = DLSAG::sign::<Sha512, OsRng>(k, ring.clone(), secret_index, &message);
            let result = DLSAG::verify::<Sha512>(signature, &message);
            assert!(result);
        }

        {
            let signature =
                DLSAG::sign::<Keccak512, OsRng>(k, ring.clone(), secret_index, &message);
            let result = DLSAG::verify::<Keccak512>(signature, &message);
            assert!(result);
        }

        {
            let signature = DLSAG::sign::<Blake2b, OsRng>(k, ring.clone(), secret_index, &message);
            let result = DLSAG::verify::<Blake2b>(signature, &message);
            assert!(result);
        }

        // Tests for signatures using the other end of the channel
        {
            let signature =
                DLSAG::sign::<Sha512, OsRng>(other_k, ring.clone(), secret_index, &message);
            let result = DLSAG::verify::<Sha512>(signature, &message);
            assert!(result);
        }

        {
            let signature =
                DLSAG::sign::<Keccak512, OsRng>(other_k, ring.clone(), secret_index, &message);
            let result = DLSAG::verify::<Keccak512>(signature, &message);
            assert!(result);
        }

        {
            let signature =
                DLSAG::sign::<Blake2b, OsRng>(other_k, ring.clone(), secret_index, &message);
            let result = DLSAG::verify::<Blake2b>(signature, &message);
            assert!(result);
        }

        let another_ring: Vec<(RistrettoPoint, RistrettoPoint, Scalar)> = (0..(n - 1)) // Prover is going to add her key into this mix
            .map(|_| {
                (
                    RistrettoPoint::random(&mut csprng),
                    RistrettoPoint::random(&mut csprng),
                    Scalar::random(&mut csprng),
                )
            })
            .collect();
        let another_message: Vec<u8> = b"This is another message".iter().cloned().collect();
        let signature_1 =
            DLSAG::sign::<Blake2b, OsRng>(k, another_ring.clone(), secret_index, &another_message);
        let signature_2 = DLSAG::sign::<Blake2b, OsRng>(k, ring.clone(), secret_index, &message);
        let signature_3 =
            DLSAG::sign::<Blake2b, OsRng>(other_k, ring.clone(), secret_index, &message);
        let result_1 = DLSAG::link(signature_1.clone(), signature_2);
        assert!(result_1);
        let result_2 = DLSAG::link(signature_1.clone(), signature_3);
        assert!(result_2);
    }
}