//! # Elgamal Public Key Encryption
//!
//! This document represents an executable specification of the Elgamal Public Key Encryption scheme.
//!
//! At the basis of the encryption scheme is a prime order group $`\mathbb{G}`$ with the same interface as defined in [draft-oprf].
#![warn(missing_docs)]

use hacspec_lib::Randomness;
use p256::{p256_point_mul, p256_point_mul_base, point_add, random_scalar, P256Point, P256Scalar};

use std::ops::Neg;

/// Domain separator for scalar sampling during key generation.
const DST_KEYGEN: &[u8] = b"ElGamal-KeyGeneration";
/// Domain separator for scalar sampling during encryption.
const DST_ENCRYPT: &[u8] = b"ElGamal-Encryption";
/// Domain separator for scalar sampling during ciphertext rerandomization.
const DST_RERANDOMIZE: &[u8] = b"ElGamal-Rerandomize";

/// An Elgamal encryption of a message `M` under public key `PK` is a pair of group elements `(c0, c1)` where
/// * `c0` is the group generator multiplied by a randomizer `r`
/// * `c1` is the result of multiplying the  target public encryption key the same randomizer `r` and adding the result to the message `M` that is to be encrypted.
pub type Ciphertext = (P256Point, P256Point);

/// An Elgamal plaintext is a member of the base group.
pub type Plaintext = P256Point;

#[derive(Debug)]
/// Any of the given algorithms may fail if the given arguments result in invalid group operations.
pub enum Error {
    /// Wrapper around Errors given by the underlying group implementation.
    CurveError,
    /// Wrapper around rejection sampling Errors
    SamplingError,
}

impl From<p256::Error> for Error {
    fn from(_value: p256::Error) -> Self {
        Self::CurveError
    }
}

/// An Elgamal private decryption key is a scalar $`dk`$ of base group $`\mathcal{G}`$.
pub type DecryptionKey = P256Scalar;
/// An Elgamal encryption key is a group element in $`\mathcal{G}$.
pub type EncryptionKey = P256Point;

/// To generate a pair of encryption and decryption keys, first the decryption key is drawn uniformly at random from the set of scalars.
/// Then the corresponding encryption key is generated by scalar multiplication of the group generator with the decryption key.
pub fn generate_keys(randomness: &mut Randomness) -> Result<(DecryptionKey, EncryptionKey), Error> {
    let dec_key = random_scalar(randomness, DST_KEYGEN)?;
    let enc_key = p256_point_mul_base(dec_key)?;

    Ok((dec_key, enc_key))
}

/// To encrypt a `message` (a group element) under encryption key `ek`, the encryption key randomized by scalar muliplication with `randomizer`. The result is used as a blinding element by adding it to the message using the group operation.
/// To allow for decryption, the ciphertext also includes as an auxillary component the result of scalar multiplication of the generator by the `randomizer` that was used in the blinding.
///

pub fn encrypt(
    enc_key: P256Point,
    message: Plaintext,
    randomness: &mut Randomness,
) -> Result<Ciphertext, Error> {
    let randomizer = random_scalar(randomness, DST_ENCRYPT)?;

    let blinded_message = p256::point_add(message, p256::p256_point_mul(randomizer, enc_key)?)?;
    let auxillary = p256::p256_point_mul_base(randomizer)?;
    Ok((auxillary, blinded_message))
}

/// To decrypt an Elgamal ciphertext, the holder of the decryption key can multiply the auxillary component of the ciphertext by the decryption key, thus reproducing the blinding element that was calculated during encryption. Adding the negation of this blinding to the blinded message component of the ciphertext recovers the encrypted message.
pub fn decrypt(dec_key: P256Scalar, ctx: Ciphertext) -> Result<Plaintext, Error> {
    let blinding: P256Point = p256::p256_point_mul(dec_key, ctx.0)?;

    p256::point_add(ctx.1, blinding.neg()).map_err(|e| e.into())
}

/// Given the correct public encryption key, it is possible to rerandomize Elgamal ciphertexts without changing the message that is encrypted.
/// To do so, the encryption is multiplied by the fresh randomizer and the result is added to the blinded message component of the ciphertext.
///
/// The auxillary element of the ciphertext is similarly updated by adding the result of multiplying the generator by the fresh randomizer.
pub fn rerandomize(
    enc_key: P256Point,
    ctx: Ciphertext,
    randomness: &mut Randomness,
) -> Result<Ciphertext, Error> {
    let randomizer = random_scalar(randomness, DST_RERANDOMIZE)?;
    let updated_blinding = p256_point_mul(randomizer, enc_key)?;
    let updated_blinded_message = point_add(updated_blinding, ctx.1)?;

    let updated_auxillary = point_add(p256_point_mul_base(randomizer)?, ctx.0)?;
    Ok((updated_auxillary, updated_blinded_message))
}

/// Elgamal ciphertexts under the same encryption key enjoy an additive homomorphism,
///  i.e. if `c1 = Enc(ek, m1)` and `c2 = Enc(pk, m2)`, then we can define an addition operation `+` such that
/// `Dec(dk, c1 + c2) = m1 + m2`.
///
/// The addition is defined as componentwise application of the group operation on the pairs of group elements that make up ciphertexts.
pub fn add_ciphertexts(lhs: Ciphertext, rhs: Ciphertext) -> Result<Ciphertext, Error> {
    let sum_blinded_messages = p256::point_add(lhs.0, rhs.0)?;
    let sum_auxillaries = p256::point_add(lhs.1, rhs.1)?;

    Ok((sum_blinded_messages, sum_auxillaries))
}

/// Since addition as described above is defined via component-wise application of the group operation, repeated addition of a ciphertext with itself
/// can be expressed via componentwise scalar multiplication of the ciphertext.
pub fn scalar_mul_ciphertext(scalar: P256Scalar, ctx: Ciphertext) -> Result<Ciphertext, Error> {
    let multiple_blinded_message = p256_point_mul(scalar, ctx.0)?;
    let multiple_auxillary = p256_point_mul(scalar, ctx.1)?;

    Ok((multiple_blinded_message, multiple_auxillary))
}

#[cfg(test)]
mod test {
    use super::*;
    use p256::random_scalar;
    use rand::RngCore;

    #[test]
    fn test_correctness() {
        let msg = random_element().unwrap();
        let mut uniform_bytes = vec![0u8; 32];
        rand::thread_rng().fill_bytes(&mut uniform_bytes);
        let mut uniform_bytes = Randomness::new(uniform_bytes);
        let (dk, ek) = generate_keys(&mut uniform_bytes).unwrap();

        let ctx = encrypt(ek, msg, &mut Randomness::new(vec![0xab; 32])).unwrap();
        let decryption = decrypt(dk, ctx).unwrap();

        debug_assert_eq!(msg, decryption);
    }

    #[test]
    fn test_rerandomize() {
        let msg = random_element().unwrap();

        let mut uniform_bytes = vec![0u8; 32];
        rand::thread_rng().fill_bytes(&mut uniform_bytes);
        let mut uniform_bytes = Randomness::new(uniform_bytes);
        let (dk, ek) = generate_keys(&mut uniform_bytes).unwrap();

        let ctx = encrypt(ek, msg, &mut Randomness::new(vec![0xab; 32])).unwrap();
        let rctx = rerandomize(ek, ctx, &mut Randomness::new(vec![0xab; 32])).unwrap();
        let decryption = decrypt(dk, rctx).unwrap();

        debug_assert_eq!(msg, decryption);
        assert_ne!(ctx, rctx);
    }

    fn random_element() -> Result<P256Point, Error> {
        let rand = random_scalar(&mut Randomness::new(vec![0xab; 32]), b"ElGamal-Test").unwrap();
        let res = p256_point_mul_base(rand)?;

        Ok(res)
    }
}