#![deny(missing_docs)]
#![feature(array_methods)]
#![feature(external_doc)]
#![doc(include = "../README.md")]
#![doc(include = "../ANONYMITY.md")]
use bit_vec::BitVec;
use curve25519_dalek::constants::RISTRETTO_BASEPOINT_POINT;
use curve25519_dalek::digest::Digest;
use curve25519_dalek::ristretto::RistrettoPoint;
use curve25519_dalek::scalar::Scalar;
use curve25519_dalek::traits::MultiscalarMul;
use rand::rngs::OsRng;
use serde::Deserialize;
use serde::Serialize;
use sha3::Sha3_512;
use std::fmt;
use std::fmt::{Display, Formatter};
use std::ops::{Mul, Sub};

#[cfg(feature = "entangled")]
use std::sync::Arc;

#[cfg(feature = "entangled")]
use rayon::iter::ParallelIterator;
#[cfg(feature = "entangled")]
use rayon::prelude::IntoParallelIterator;

/// A tag is a probabilistic cryptographic structure. When constructed for a given `FuzzyPublicKey`
/// it will pass the `FuzzyDetectionKey::test` 100% of the time. For other public keys
/// it will pass the test with probability `gamma` related to the security parameter of the system.
/// This system provides the following security properties:
/// * Correctness: Valid tags constructed for a specific public key will always validate when tested using the detection key
/// * Fuzziness: Invalid tags will produce false positives with probability _p_ related to the security property (_γ_)
/// * Security: An adversarial server with access to the detection key is unable to distinguish false positives from true positives. (Detection Ambiguity)
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct FuzzyTag {
    u: RistrettoPoint,
    y: Scalar,
    ciphertexts: BitVec,
}

impl Display for FuzzyTag {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "{} {} {}",
            hex::encode(self.u.compress().as_bytes()),
            hex::encode(self.y.as_bytes()),
            hex::encode(self.ciphertexts.to_bytes())
        )
    }
}

/// The complete secret key. Can't directly be used for testing. Instead you will need to generate
/// a FuzzyDetectionKey using extract
#[derive(Debug, Serialize, Deserialize)]
pub struct FuzzySecretKey {
    /// the detection key - this can be given to adversarial servers to help probabilistically
    /// filter messages (with a false-positive rate derived from γ and a 0% false negative rate)
    secret_key: Vec<Scalar>,
    /// the public key - this can be given to people who you want to contact you
    public_key: FuzzyPublicKey,
}

impl FuzzySecretKey {
    /// Generate a new Key Pair given a security parameter `gamma`. Tags generated for a given
    /// `FuzzyPublicKey::flag` will pass the `FuzzyDetectionKey::test` for other public
    /// keys with probability $ 2 ^ -8 $
    /// Example:
    /// ```
    ///     use fuzzytags::{FuzzySecretKey};
    ///     let secret_key = FuzzySecretKey::generate(5);
    /// ```
    pub fn generate(gamma: usize) -> FuzzySecretKey {
        let mut rng = OsRng::default();
        let g = RISTRETTO_BASEPOINT_POINT;
        let mut secret_key = vec![];
        let mut p_keys = vec![];
        for _i in 0..gamma {
            let sk_i = Scalar::random(&mut rng);
            let pk_i = g.mul(sk_i);
            secret_key.push(sk_i);
            p_keys.push(pk_i);
        }
        FuzzySecretKey {
            secret_key,
            public_key: FuzzyPublicKey { 0: p_keys },
        }
    }

    /// extract a detection key for a given false positive (p = 2^-n)
    /// Example:
    /// ```
    ///     use fuzzytags::{FuzzySecretKey};
    ///     let secret_key = FuzzySecretKey::generate(24);
    ///     let detection_key = secret_key.extract(2);
    /// ```
    pub fn extract(&self, n: usize) -> FuzzyDetectionKey {
        let parts = self.secret_key.iter().take(n).cloned().collect();
        FuzzyDetectionKey { 0: parts }
    }

    /// derive the public key for this key
    /// Example:
    /// ```
    ///     use fuzzytags::FuzzySecretKey;
    ///     let secret_key = FuzzySecretKey::generate(24);
    ///     let public_key = secret_key.public_key();
    /// ```
    pub fn public_key(&self) -> FuzzyPublicKey {
        self.public_key.clone()
    }

    /// a hash function that takes 3 ristretto points as a parameter and outputs 0 or 1.
    fn h(u: RistrettoPoint, h: RistrettoPoint, w: RistrettoPoint) -> u8 {
        let mut hash = sha3::Sha3_256::new();
        hash.update(u.compress().as_bytes());
        hash.update(h.compress().as_bytes());
        hash.update(w.compress().as_bytes());
        return hash.finalize().as_slice()[0] & 0x01;
    }

    /// a hash function which takes a ristretto point and a vector of ciphertexts and outputs a
    /// ristretto scalar.
    fn g(u: RistrettoPoint, points: &BitVec) -> Scalar {
        let mut input = points.to_bytes().as_slice().to_vec();
        input.extend_from_slice(u.compress().as_bytes());
        Scalar::hash_from_bytes::<Sha3_512>(input.as_slice())
    }
}

/// A collection of "secret" data that can be used to determine if a `FuzzyTag` was intended for
/// the derived public key with probability p
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct FuzzyDetectionKey(Vec<Scalar>);

impl FuzzyDetectionKey {
    /// a convenient id for a detection key for internal accounting purposes
    /// do not expose this to applications
    pub fn id(&self) -> String {
        let mut hash = sha3::Sha3_256::new();
        for s in self.0.iter() {
            hash.update(s.as_bytes())
        }
        format!("{}", hex::encode(hash.finalize().as_slice()),)
    }
}

impl Display for FuzzyDetectionKey {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.id())
    }
}

impl FuzzyDetectionKey {
    /// calculate the ideal false positive rate of this detection key
    /// ```
    ///    use fuzzytags::FuzzySecretKey;
    ///    let gamma = 24;
    ///    let secret_key = FuzzySecretKey::generate(gamma);
    ///    let public_key = secret_key.public_key();
    ///    // extract a detection key
    ///    let detection_key = secret_key.extract(5);
    ///    detection_key.false_positive_probability();
    /// ```
    pub fn false_positive_probability(&self) -> f64 {
        (2.0_f64).powi(0 - (self.0.len() as i32))
    }

    /// returns true if the tag was intended for this key
    /// Example:
    /// ```
    ///    use fuzzytags::FuzzySecretKey;
    ///     let gamma = 24;
    ///    let secret_key = FuzzySecretKey::generate(gamma);
    ///    let public_key = secret_key.public_key();
    ///    // extract a detection key
    ///    let detection_key = secret_key.extract(5);
    ///
    ///    // Give public key to a another party...
    ///    // and then they can do...
    ///    let tag = public_key.generate_tag();
    ///
    ///    // The server can now do this:
    ///    if detection_key.test_tag(&tag) {
    ///    // the message attached to this tag *might* be for the party associated with the detection key
    ///    } else {
    ///    // the message attached to this tag is definitely *not* for the party associated with the detection key.
    ///    }
    /// ```
    pub fn test_tag(&self, tag: &FuzzyTag) -> bool {
        let m = FuzzySecretKey::g(tag.u, &tag.ciphertexts);

        let g = RISTRETTO_BASEPOINT_POINT;

        // Re-derive w = g^z from the public tag.
        // y = (1/r * (z-m)
        // u = g^r
        // so   w = g^m + u^y
        //      w = g^m + g^(r * 1/r * (z-m))
        //      w = g^m + g^(z-m)
        //      w = g^z
        // See below for a full explanation as to the reason for this:
        let w = RistrettoPoint::multiscalar_mul(&[m, tag.y], &[g, tag.u]);

        // for each secret key part...
        let mut result = true;
        for (i, x_i) in self.0.iter().enumerate() {
            // re-derive the key from the tag
            let k_i = FuzzySecretKey::h(tag.u, tag.u.mul(x_i), w);

            // calculate the "original" plaintext
            let c_i = match tag.ciphertexts.get(i) {
                Some(true) => 0x01,
                Some(false) => 0x00,
                _ => 0x00,
                // we've run our of ciphertext, it doesn't really matter what we put here, the rest of the test will fail
                // since the security of k_i is modelled as a random oracle, (k_i ^ 0) should also be random
            };

            let b_i = k_i ^ c_i;

            // assert that the plaintext is all 1's
            result = result & (b_i == 1);
        }
        return result;
    }
}

/// A public identity that others can create tags for.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct FuzzyPublicKey(Vec<RistrettoPoint>);

impl FuzzyPublicKey {
    /// a convenient id for a public key for internal accounting purposes
    /// do not expose this to applications
    pub fn id(&self) -> String {
        let mut hash = sha3::Sha3_256::new();
        for s in self.0.iter() {
            hash.update(s.compress().as_bytes())
        }
        format!("{}", hex::encode(hash.finalize().as_slice()),)
    }

    /// generate a new tag for this public key
    /// Example:
    /// ```
    ///     use fuzzytags::{FuzzySecretKey};
    ///     let secret_key = FuzzySecretKey::generate(24);
    ///     let public_key = secret_key.public_key(); // give this to a sender
    ///     let tag = public_key.generate_tag();
    /// ```
    pub fn generate_tag(&self) -> FuzzyTag {
        let mut rng = OsRng::default();
        let g = RISTRETTO_BASEPOINT_POINT;

        // generate some random points...
        let r = Scalar::random(&mut rng);
        let u = g.mul(r);
        let z = Scalar::random(&mut rng);
        let w = g.mul(z);

        // construct the ciphertext portion of the tag
        let mut ciphertexts = BitVec::new();
        for (_i, h_i) in self.0.iter().enumerate() {
            let k_i = FuzzySecretKey::h(u, h_i.mul(r), w);
            // encrypt a plaintext of all 1's
            let c_i = k_i ^ 0x01;
            ciphertexts.push(c_i == 0x01);
        }

        // Without this next part, this scheme would not be CCA-secure. Consider a scheme with just
        // u = ^r and and h_i^r = g^(x_i*r)
        // An adversarial server with access to a Test oracle (i.e. the decryption key) may be able
        // to maul a challenge ciphertext by e.g. replacing the order of the ciphertexts.

        // From the paper:
        // "The value w corresponds to a chameleon hash [KR00] computed on the message (0,z), where z is chosen at random.
        // Once the ciphertext has been computed, we use a master trapdoor for the chameleon hash (which is part of the scheme’s secret key) in order to compute a collision (y,m) where m
        // is a hash of the remaining components of the ciphertext"

        // Translated m is a challenge over the random element u and the ordered ciphertexts
        // It is then used to construct a response y which can be used to recover w the random element
        // used to derive the key.

        // finally calculate a `y` = 1/r * (z-m) which will be used to re-derive `w`
        let m = FuzzySecretKey::g(u, &ciphertexts);
        let y = r.invert().mul(z.sub(m));

        return FuzzyTag { u, y, ciphertexts };
    }

    #[cfg(feature = "entangled")]
    /// WARNING: if you pass in a large length into this function it will take a long time!
    /// This begins a very slow, but ,parallel search for a tag that will validate of the given
    /// public keys up to a given false positive rate 2^-1
    /// Example:
    /// ```
    ///     use fuzzytags::{FuzzySecretKey, FuzzyPublicKey};
    ///     let secret_key_1 = FuzzySecretKey::generate(24);
    ///     let secret_key_2 = FuzzySecretKey::generate(24);
    ///     let public_key_1 = secret_key_1.public_key(); // give this to a sender
    ///     let public_key_2 = secret_key_2.public_key(); // give this to a sender
    ///     // Will validate for detection keys derived from both secret_key_1 and secret_key_2 up
    ///     // to n=8
    ///     let tag = FuzzyPublicKey::generate_entangled_tag(vec![public_key_1,public_key_2], 8);
    /// ```
    pub fn generate_entangled_tag(public_keys: Vec<FuzzyPublicKey>, length: usize) -> FuzzyTag {
        let arc_public_keys = Arc::new(public_keys);
        loop {
            let results: Vec<FuzzyTag> = (0..8)
                .into_par_iter()
                .map(|_x| FuzzyPublicKey::try_entangled_tag(arc_public_keys.clone(), length))
                .filter(|x| x.is_ok())
                .map(|x| x.unwrap())
                .collect();
            if results.is_empty() == false {
                return results[0].clone();
            }
        }
    }

    #[cfg(feature = "entangled")]
    fn try_entangled_tag(public_keys: Arc<Vec<FuzzyPublicKey>>, length: usize) -> Result<FuzzyTag, ()> {
        let mut rng = OsRng::default();
        let g = RISTRETTO_BASEPOINT_POINT;

        // generate some random points...
        let r = Scalar::random(&mut rng);
        let u = g.mul(r);

        let mut entangled = false;
        let mut z = Scalar::zero();

        // construct the ciphertext portion of the tag
        let mut ciphertexts = BitVec::new();
        let mut attempts = 0;

        let mut public_key_precomputes = vec![];
        for public_key in public_keys.iter() {
            let mut precompute = vec![];
            for i in public_key.0.iter() {
                precompute.push(i.mul(r));
            }
            public_key_precomputes.push(precompute);
        }

        while !entangled && attempts < 1000 {
            attempts += 1;
            ciphertexts = BitVec::new();
            z = Scalar::random(&mut rng);
            let w = g.mul(z);
            entangled = true;
            for (i, precompute) in public_key_precomputes[0].iter().enumerate() {
                let mut same = true;
                let k_i = FuzzySecretKey::h(u, *precompute, w);

                if i < length {
                    for precompute in public_key_precomputes.iter().skip(1) {
                        let n_k_i = FuzzySecretKey::h(u, precompute[i], w);
                        if k_i != n_k_i {
                            same = false;
                            break;
                        }
                    }
                    if !same {
                        entangled = false;
                        break;
                    }
                }
                // encrypt a plaintext of all 1's
                let c_i = k_i ^ 0x01;
                ciphertexts.push(c_i == 0x01);
            }
        }

        if entangled == false {
            return Err(());
        }

        // Without this next part, this scheme would not be CCA-secure. Consider a scheme with just
        // u = ^r and and h_i^r = g^(x_i*r)
        // An adversarial server with access to a Test oracle (i.e. the decryption key) may be able
        // to maul a challenge ciphertext by e.g. replacing the order of the ciphertexts.

        // From the paper:
        // "The value w corresponds to a chameleon hash [KR00] computed on the message (0,z), where z is chosen at random.
        // Once the ciphertext has been computed, we use a master trapdoor for the chameleon hash (which is part of the scheme’s secret key) in order to compute a collision (y,m) where m
        // is a hash of the remaining components of the ciphertext"

        // Translated m is a challenge over the random element u and the ordered ciphertexts
        // It is then used to construct a response y which can be used to recover w the random element
        // used to derive the key.

        // finally calculate a `y` = 1/r * (z-m) which will be used to re-derive `w`
        let m = FuzzySecretKey::g(u, &ciphertexts);
        let y = r.invert().mul(z.sub(m));

        return Ok(FuzzyTag { u, y, ciphertexts });
    }
}

#[cfg(test)]
mod tests {
    use crate::FuzzySecretKey;

    #[test]
    fn test_serialization() {
        let secret_key = FuzzySecretKey::generate(24);
        let tag = secret_key.public_key.generate_tag();
        let detection_key = secret_key.extract(10);
        println!("{}", serde_json::to_string(&tag).unwrap());
        println!("{}", serde_json::to_string(&detection_key).unwrap());
    }

    #[test]
    #[cfg(feature = "entangled")]
    fn test_multiple() {
        use crate::FuzzyPublicKey;
        let secret_keys: Vec<FuzzySecretKey> = (0..3).map(|_x| FuzzySecretKey::generate(24)).collect();
        let public_keys: Vec<FuzzyPublicKey> = secret_keys.iter().map(|x| x.public_key()).collect();
        // it takes ~15 minutes on a standard desktop to find a length=24 match for 2 parties, so for testing let's keep things light
        let entangled_tag = FuzzyPublicKey::generate_entangled_tag(public_keys, 6);
        println!("{}", entangled_tag);
        for secret_key in secret_keys.iter() {
            let detection_key = secret_key.extract(6);
            assert!(detection_key.test_tag(&entangled_tag));
            println!("{}", detection_key);
        }
    }

    #[test]
    fn correctness() {
        let number_of_messages = 100;
        let secret_key = FuzzySecretKey::generate(16);
        for i in 0..number_of_messages {
            let tag = secret_key.public_key().generate_tag();
            println!("{}: {}", i, tag);
            assert!(secret_key.extract(5).test_tag(&tag));
        }
    }

    #[test]
    fn false_positives() {
        let gamma = 8;
        let number_of_messages = 1000;
        let secret_key = FuzzySecretKey::generate(gamma);
        let mut false_positives = 0;
        for _i in 0..number_of_messages {
            let secret_key2 = FuzzySecretKey::generate(gamma);
            let tag = secret_key2.public_key().generate_tag();
            assert!(secret_key2.extract(3).test_tag(&tag));
            if secret_key.extract(3).test_tag(&tag) == true {
                false_positives += 1;
            }
        }
        println!(
            "Expected False Positive Rate: {}\nActual False Positive Rate: {}",
            secret_key.extract(3).false_positive_probability(),
            (false_positives as f64 / number_of_messages as f64)
        );
    }
}