#![deny(missing_docs)]
#![feature(external_doc)]
#![doc(include = "../README.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 rand::rngs::OsRng;
use serde::Deserialize;
use serde::Serialize;
use sha3::Sha3_512;
use std::fmt;
use std::fmt::{Display, Formatter};
use std::ops::{Add, Mul, Sub};

/// 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(Debug, Serialize, Deserialize)]
pub struct FuzzyTag {
    u: RistrettoPoint,
    y: Scalar,
    ciphertexts: BitVec,
}

/// The complete secret key. Can't directly be used for testing. Instead you will need to generate
/// a FuzzyDetectionKey using extract
pub struct FuzzySecretKey(Vec<Scalar>);

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

/// 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(Debug, Serialize, Deserialize)]
pub struct FuzzyDetectionKey(Vec<Scalar>);

impl FuzzyDetectionKey {
    /// returns true if the tag was intended for this key
    pub fn test_tag(&self, tag: &FuzzyTag) -> bool {
        let m = FuzzyTagKeyPair::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 = g.mul(m).add(tag.u.mul(tag.y));

        // 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 = FuzzyTagKeyPair::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.
pub struct FuzzyPublicKey(Vec<RistrettoPoint>);

impl FuzzyPublicKey {
    /// generate a new tag for this public key
    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 = FuzzyTagKeyPair::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 = FuzzyTagKeyPair::g(u, &ciphertexts);
        let y = r.invert().mul(z.sub(m));

        return FuzzyTag { u, y, ciphertexts };
    }
}

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())
        )
    }
}

/// An identity keypair for generating and validating fuzzy meta tags.
pub struct FuzzyTagKeyPair {
    /// 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)
    pub secret_key: FuzzySecretKey,
    /// the public key - this can be given to people who you want to contact you
    pub public_key: FuzzyPublicKey,
}

impl FuzzyTagKeyPair {
    /// 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 $
    pub fn generate(gamma: usize) -> FuzzyTagKeyPair {
        let mut rng = OsRng::default();
        let g = RISTRETTO_BASEPOINT_POINT;
        let mut s_keys = 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);
            s_keys.push(sk_i);
            p_keys.push(pk_i);
        }
        FuzzyTagKeyPair {
            secret_key: FuzzySecretKey { 0: s_keys },
            public_key: FuzzyPublicKey { 0: p_keys },
        }
    }

    /// extract a detection key for a given false positive (p = 2^-n)
    /// a facade for an extraction of the encapsulated secret key
    pub fn extract(&self, n: usize) -> FuzzyDetectionKey {
        self.secret_key.extract(n)
    }

    /// a hash function that takes 3 risretto points as a parameter and outputs 0 or 1.
    fn h(u: RistrettoPoint, h: RistrettoPoint, w: RistrettoPoint) -> u8 {
        let hash = sha3::Sha3_256::digest(
            format!(
                "{}{}{}",
                hex::encode(u.compress().as_bytes()),
                hex::encode(h.compress().as_bytes()),
                hex::encode(w.compress().as_bytes())
            )
            .as_bytes(),
        );
        return hash.as_slice()[0] & 0x01;
    }

    /// a hash function which takes a ristretto point and a vector of ciphertexts and ouputs a
    /// ristretto scalar.
    fn g(u: RistrettoPoint, points: &BitVec) -> Scalar {
        Scalar::hash_from_bytes::<Sha3_512>(format!("{}{}", hex::encode(u.compress().as_bytes()), hex::encode(points.to_bytes())).as_bytes())
    }
}

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

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

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

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