// Copyright 2019 Stichting Organism
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Elgamal encryption

use mohan::{
    ser,
    dalek::{
        scalar::Scalar,
        constants::RISTRETTO_BASEPOINT_TABLE,
        ristretto::CompressedRistretto,
        traits::Identity
    },
    mohan_rand
};
use bacteria::Transcript;
use schnorr::SecretKey;
use schnorr::PublicKey;
use crate::ZeiError;
use serde::{ Serialize, Deserialize };


#[derive(Clone, PartialEq, Eq, Debug, Serialize, Deserialize)]
/// On Chain Magic Box
pub struct Lockbox {
    /// our blinded randomned g^R
    pub rand: CompressedRistretto,
    /// Encrypted data
    pub data: Vec<u8>,
}

impl Default for Lockbox {
    fn default() -> Lockbox {
        Lockbox {
            rand: CompressedRistretto::identity(),
            data: Vec::new()
        }
    }
}

//locks a given amount with reciever PK and returns cipher text and rand
pub fn lock(randomness: &Scalar, publickey: &PublicKey, message: &[u8]) -> Lockbox {
        //New strobe context
        let mut ctx = Transcript::new(b"Zei.Lockbox");
        //commit pk
        ctx.append_message(b"PK", publickey.as_bytes());
        
        //get our enc key and blinded randomness
        //g^(x*r), g^r
        
        //pk^R
        let shared_key = randomness * publickey.as_point();

        //g^R where R = randomness used to derive shared key
        let blind_rand = randomness * &RISTRETTO_BASEPOINT_TABLE;
        
        //key strobe with shared key
        ctx.strobe.key(shared_key.compress().as_bytes(), false);

        let mut buf = message.to_vec();
        ctx.strobe.send_enc(buf.as_mut_slice(), false);
        
        return Lockbox {
            data: buf,
            rand: blind_rand.compress()
        };
}


/// Open the box with the corresponding secret key & randomness
pub fn unlock(boxy: &mut Lockbox, secretkey: &SecretKey, publickey: &PublicKey) -> Result<(), ZeiError> {
    //g^x = pk. x is secret key, g^R == rand 
    //New strobe context
    let mut ctx = Transcript::new(b"Zei.Lockbox");

    //commit pk
    ctx.append_message(b"PK", publickey.as_bytes());

    //REDERIVED_KEY == (g^R)^x
    let dec_key = boxy.rand.decompress().unwrap() * secretkey.as_scalar();
    
    //feed our secret
    ctx.strobe.key(dec_key.compress().as_bytes(), false);
    //open sesame
    ctx.strobe.recv_enc(&mut boxy.data, false);

    Ok(())
    
}

impl ser::Writeable for Lockbox {
	fn write<W: ser::Writer>(&self, writer: &mut W) -> Result<(), ser::Error> {
		self.rand.write(writer)?;
        //Data length
        writer.write_u64(self.data.len() as u64)?;
        //Raw Data
		self.data.write(writer)?;

        Ok(())
	}
}

impl ser::Readable for Lockbox {
	fn read(reader: &mut dyn ser::Reader) -> Result<Lockbox, ser::Error> {
		let rand = CompressedRistretto::read(reader)?;
        let data_len = reader.read_u64()?;
        let raw_data = reader.read_fixed_bytes(data_len as usize)?;
      
		Ok(Lockbox{
            rand: rand,
            data: raw_data
        })
	}
}


#[cfg(test)]
mod test {
    use super::*;
    use mohan::mohan_rand;
    use schnorr::Keypair;
    use mohan::byteorder::{ByteOrder, BigEndian};


    #[test]
    fn test_lock_unlock() { 
        //Sample Fresh Keypair
        let mut csprng = mohan_rand();
        let keypair: Keypair = Keypair::generate(&mut csprng);
        
        //1. Sample Fresh blinding factor [blind], its a scalar
        let blinding_t = Scalar::random(&mut csprng);
        // coins to send
        let amount: u32 = 101;

        //7. Encrypt to receiver pubkey both the amount transferred and the blinding factor [blind] 
        let mut to_encrypt = Vec::new();
        //first add amount which is fixed 4 bytes in big endian
        let mut encoded_amount = [0; 4];
        BigEndian::write_u32(&mut encoded_amount, amount);
       
        //println!("encoded_amount: {:?}", encoded_amount);
        to_encrypt.extend_from_slice(&encoded_amount);
        //next add the blind
        to_encrypt.extend_from_slice(&blinding_t.to_bytes());

        //println!("to_encrypt: {:?}", to_encrypt);
        
        //lock em up
        let rand_t = Scalar::random(&mut csprng);
        let mut lbox = lock(&rand_t, &keypair.public, &to_encrypt);

        //Now we unbox to check if we get same results
        
        //unlock encrypted box
        let unlocked = unlock(&mut lbox, &keypair.secret, &keypair.public);
        assert!(unlocked.is_ok());
        //extract balance value & blind value
        let (raw_amount, _raw_blind) = lbox.data.split_at(5);
        //println!("unlocked value: {:?}", unlocked);
        //convert to u32
        let p_amount = BigEndian::read_u32(&raw_amount);
        
        //println!("amount open: {:?}", p_amount);
        //check if amounts are the same
        assert_eq!(p_amount, amount);
    }

     #[test]
    fn test_box_roundtrip() { 
        //Sample Fresh Keypair
         let mut csprng = mohan_rand();
        let keypair: Keypair = Keypair::generate(&mut csprng);
        
        //1. Sample Fresh blinding factor [blind], its a scalar
        let blinding_t = Scalar::random(&mut csprng);
        // coins to send
        let amount: u32 = 101;

        //7. Encrypt to receiver pubkey both the amount transferred and the blinding factor [blind] 
        let mut to_encrypt = Vec::new();
        //first add amount which is fixed 4 bytes in big endian
        let mut encoded_amount = [0; 4];
        BigEndian::write_u32(&mut encoded_amount, amount);
       
        //println!("encoded_amount: {:?}", encoded_amount);
        to_encrypt.extend_from_slice(&encoded_amount);
        //next add the blind
        to_encrypt.extend_from_slice(&blinding_t.to_bytes());

        //println!("to_encrypt: {:?}", to_encrypt);
        
        //lock em up
        let rand_t = Scalar::random(&mut csprng);
        let lbox = lock(&rand_t, &keypair.public, &to_encrypt);

        //Now we ser/der
        use std::io::{self, Write, Read};

        // Write
        let mut tmpfile = tempfile::NamedTempFile::new().unwrap();
        // Re-open it.
        let mut read_handle = tmpfile.reopen().unwrap();

        let _x = mohan::ser::serialize_default(&mut tmpfile, &lbox).expect("serialization failed");
        let mut lbox: Lockbox = mohan::ser::deserialize_default(&mut read_handle).unwrap();
            
        //unlock encrypted box
        let unlocked = unlock(&mut lbox, &keypair.secret, &keypair.public);
        assert!(unlocked.is_ok());
        //extract balance value & blind value
        let (raw_amount, _raw_blind) = lbox.data.split_at(5);
        //println!("unlocked value: {:?}", unlocked);
        //convert to u32
        let p_amount = BigEndian::read_u32(&raw_amount);
        
        //println!("amount open: {:?}", p_amount);
        //check if amounts are the same
        assert_eq!(p_amount, amount);
    }

}