extern crate paillier;

extern crate serde_json;
extern crate rand;

use paillier::*;

#[cfg(not(feature="keygen"))]
fn main() {
    println!("*** please run with 'keygen' feature ***")
}

#[cfg(feature="keygen")]
fn main() {

    // first we initialize the clerk by asking it to generate a fresh keypair
    let clerk = Clerk::new();

    // then we create a set of voters using the corresponding encryption key of the clerk
    let ek = clerk.encryption_key();
    let voters = (0..10)
        .map(|_| Voter::new(&ek) )
        .collect::<Vec<_>>();

    // the clerk launched a new vote by sharing an encryption of zero
    let mut tally = clerk.new_voting();

    // each voter in turn
    voters.iter().for_each(|voter| {
        tally = voter.vote(&tally);
    });

    // let clerk reveal final tally
    let nb_voters_for = clerk.reveal(&tally);
    let nb_voters_against = voters.len() as u64 - nb_voters_for;
    println!("The result is {} for and {} against", nb_voters_for, nb_voters_against);
    
    // check tally correctness for fun 
    // - normally this wouldn't be possible of course
    assert_eq!(nb_voters_for, voters.iter().filter(|voter| voter.vote).count() as u64);
    assert_eq!(nb_voters_against, voters.iter().filter(|voter| !voter.vote).count() as u64);
}

struct Clerk { 
    ek: EncryptionKey,
    dk: DecryptionKey,
}

impl Clerk {

    fn new() -> Self {
        // generate fresh keypair
        let keypair = Paillier::keypair();
        // extract encryption key from keypair
        let (ek, dk) = keypair.keys();
        Clerk { ek, dk }
    }

    fn encryption_key(&self) -> String {
        // serialize key for sending
        serde_json::to_string(&self.ek).unwrap()
    }

    fn new_voting(&self) -> String {
        // encrypt zero
        let c = Paillier::encrypt(&self.ek, 0);
        // serialize the ciphertext
        serde_json::to_string(&c).unwrap()
    }

    fn reveal(&self, tally: &str) -> u64 {
        // deserialize ciphertext
        let c: EncodedCiphertext<u64> = serde_json::from_str(tally).unwrap();
        // decrypt tally
        Paillier::decrypt(&self.dk, c)
    }

}

struct Voter {
    ek: EncryptionKey,
    vote: bool,
}

impl Voter {

    fn new(ek: &str) -> Voter {
        // deserialize encryption key
        let ek: EncryptionKey = serde_json::from_str(&ek).unwrap();
        // generate random vote
        use rand::Rng;
        let mut rng = rand::thread_rng();
        let vote = rng.gen();

        Voter { ek, vote }
    }

    fn vote(&self, tally: &str) -> String {
        // deserialize current tally ciphertext
        let c: EncodedCiphertext<u64> = serde_json::from_str(tally).unwrap();
        // add own vote
        let d = Paillier::add(&self.ek, c, if self.vote {1} else {0});
        // re-randomize once all homomorphic operations have been performed
        let d = Paillier::rerandomize(&self.ek, d);
        // re-serialize ciphertext
        serde_json::to_string(&d).unwrap()
    }

}