//! Arcis implementation of <https://github.com/solana-program/zk-elgamal-proof/blob/main/zk-sdk/src/sigma_proofs/zero_ciphertext.rs>

use crate::{
    core::{
        circuits::boolean::{boolean_value::BooleanValue, byte::Byte},
        global_value::{curve_value::CompressedCurveValue, value::FieldValue},
    },
    traits::{Reveal, ToLeBytes},
    utils::{
        field::ScalarField,
        zkp::{
            elgamal::{ElGamalCiphertext, ElGamalKeypair, ElGamalPubkey},
            transcript::Transcript,
            util::UNIT_LEN,
        },
    },
};
use zk_elgamal_proof::encryption::{
    DECRYPT_HANDLE_LEN,
    ELGAMAL_CIPHERTEXT_LEN,
    ELGAMAL_PUBKEY_LEN,
    PEDERSEN_COMMITMENT_LEN,
};

pub const ZERO_CIPHERTEXT_PROOF_LEN: usize = 96;

pub const ZERO_CIPHERTEXT_PROOF_CONTEXT_LEN: usize = ELGAMAL_PUBKEY_LEN + ELGAMAL_CIPHERTEXT_LEN;

pub const ZERO_CIPHERTEXT_PROOF_DATA_LEN: usize =
    ZERO_CIPHERTEXT_PROOF_CONTEXT_LEN + ZERO_CIPHERTEXT_PROOF_LEN;

/// The instruction data that is needed for the `ProofInstruction::ZeroCiphertext` instruction.
///
/// It includes the cryptographic proof as well as the context data information needed to verify
/// the proof.
#[derive(Clone, Copy)]
pub struct ZeroCiphertextProofData {
    pub context: ZeroCiphertextProofContext,
    pub proof: ZeroCiphertextProof,
}

/// The context data needed to verify a zero-ciphertext proof.
#[derive(Clone, Copy)]
pub struct ZeroCiphertextProofContext {
    /// The ElGamal pubkey associated with the ElGamal ciphertext
    pub pubkey: ElGamalPubkey,
    /// The ElGamal ciphertext that encrypts zero
    pub ciphertext: ElGamalCiphertext,
}

impl ZeroCiphertextProofContext {
    fn new_transcript(&self) -> Transcript<BooleanValue> {
        let mut transcript = Transcript::new(b"zero-ciphertext-instruction");
        transcript.append_point(b"pubkey", &self.pubkey.get_point().compress());
        transcript.append_elgamal_ciphertext(b"ciphertext", &self.ciphertext);
        transcript
    }

    pub fn to_bytes(&self) -> [Byte<BooleanValue>; ZERO_CIPHERTEXT_PROOF_CONTEXT_LEN] {
        let mut bytes = [Byte::<BooleanValue>::from(0); ZERO_CIPHERTEXT_PROOF_CONTEXT_LEN];
        bytes[..ELGAMAL_PUBKEY_LEN].copy_from_slice(&self.pubkey.get_point().compress().to_bytes());
        let mut offset = ELGAMAL_PUBKEY_LEN;
        bytes[offset..offset + PEDERSEN_COMMITMENT_LEN]
            .copy_from_slice(&self.ciphertext.commitment.get_point().compress().to_bytes());
        offset += PEDERSEN_COMMITMENT_LEN;
        bytes[offset..offset + DECRYPT_HANDLE_LEN]
            .copy_from_slice(&self.ciphertext.handle.get_point().compress().to_bytes());
        bytes
    }
}

impl ZeroCiphertextProofData {
    pub fn new(keypair: &ElGamalKeypair, ciphertext: &ElGamalCiphertext) -> Self {
        let context = ZeroCiphertextProofContext {
            pubkey: *keypair.pubkey(),
            ciphertext: *ciphertext,
        };
        let mut transcript = context.new_transcript();
        let proof = ZeroCiphertextProof::new(keypair, ciphertext, &mut transcript);
        ZeroCiphertextProofData { context, proof }
    }

    pub fn to_bytes(&self) -> [Byte<BooleanValue>; ZERO_CIPHERTEXT_PROOF_DATA_LEN] {
        let mut bytes = [Byte::<BooleanValue>::from(0); ZERO_CIPHERTEXT_PROOF_DATA_LEN];
        bytes[..ZERO_CIPHERTEXT_PROOF_CONTEXT_LEN].copy_from_slice(&self.context.to_bytes());
        bytes[ZERO_CIPHERTEXT_PROOF_CONTEXT_LEN..].copy_from_slice(&self.proof.to_bytes());
        bytes
    }
}

/// Zero-ciphertext proof.
///
/// Contains all the elliptic curve and scalar components that make up the sigma protocol.
#[allow(non_snake_case, dead_code)]
#[derive(Clone, Copy)]
pub struct ZeroCiphertextProof {
    Y_P: CompressedCurveValue,
    Y_D: CompressedCurveValue,
    z: FieldValue<ScalarField>,
}

#[allow(non_snake_case)]
impl ZeroCiphertextProof {
    /// Creates a zero-ciphertext proof.
    ///
    /// The function does *not* hash the public key and ciphertext into the transcript. For
    /// security, the caller (the main protocol) should hash these public components prior to
    /// invoking this constructor.
    ///
    /// This function is randomized. It uses random singlets internally to generate random scalars.
    ///
    /// * `elgamal_keypair` - The ElGamal keypair associated with the ciphertext to be proved
    /// * `ciphertext` - The main ElGamal ciphertext to be proved
    /// * `transcript` - The transcript that does the bookkeeping for the Fiat-Shamir heuristic
    pub fn new(
        elgamal_keypair: &ElGamalKeypair,
        ciphertext: &ElGamalCiphertext,
        transcript: &mut Transcript<BooleanValue>,
    ) -> Self {
        transcript.zero_ciphertext_proof_domain_separator();

        // extract the relevant scalar and Ristretto points from the input
        let P = elgamal_keypair.pubkey().get_point();
        let s = elgamal_keypair.secret().get_scalar();
        let D = ciphertext.handle.get_point();

        // generate a random masking factor that also serves as a nonce
        let y = FieldValue::<ScalarField>::random();
        let Y_P = (y * *P).reveal().compress();
        let Y_D = (y * *D).reveal().compress();

        // record Y in the transcript and receive a challenge scalar
        transcript.append_point(b"Y_P", &Y_P);
        transcript.append_point(b"Y_D", &Y_D);

        let c = transcript.challenge_scalar(b"c");

        // compute the masked secret key
        let z = ((c * *s) + y).reveal();

        transcript.append_scalar(b"z", &z);
        let _w = transcript.challenge_scalar(b"w");

        Self { Y_P, Y_D, z }
    }

    pub fn to_bytes(&self) -> [Byte<BooleanValue>; ZERO_CIPHERTEXT_PROOF_LEN] {
        let mut buf = [Byte::<BooleanValue>::from(0); ZERO_CIPHERTEXT_PROOF_LEN];
        let mut chunks = buf.chunks_mut(UNIT_LEN);
        chunks.next().unwrap().copy_from_slice(&self.Y_P.to_bytes());
        chunks.next().unwrap().copy_from_slice(&self.Y_D.to_bytes());
        chunks
            .next()
            .unwrap()
            .copy_from_slice(&self.z.to_le_bytes());
        buf
    }
}

#[cfg(test)]
mod tests {
    use crate::{
        core::{
            bounds::FieldBounds,
            expressions::{
                curve_expr::{CurveExpr, InputInfo},
                domain::Domain,
                expr::EvalValue,
                field_expr::FieldExpr,
                InputKind,
            },
            global_value::{
                curve_value::CurveValue,
                global_expr_store::with_local_expr_store_as_global,
                value::FieldValue,
            },
            ir_builder::{ExprStore, IRBuilder},
        },
        utils::{
            curve_point::CurvePoint,
            field::ScalarField,
            zkp::{
                elgamal::{
                    DecryptHandle,
                    ElGamalCiphertext,
                    ElGamalKeypair,
                    ElGamalPubkey,
                    ElGamalSecretKey,
                },
                pedersen::PedersenCommitment,
                zero_ciphertext::ZeroCiphertextProofData,
            },
        },
    };
    use group::GroupEncoding;
    use primitives::algebra::elliptic_curve::{Curve as AsyncMPCCurve, Curve25519Ristretto};
    use rand::{Rng, RngCore};
    use rustc_hash::FxHashMap;
    use std::rc::Rc;
    use zk_elgamal_proof::{
        encryption::elgamal::ElGamalKeypair as SolanaElGamalKeypair,
        zk_elgamal_proof_program::proof_data::{
            ZeroCiphertextProofData as SolanaZeroCiphertextProofData,
            ZkProofData,
        },
    };

    #[test]
    #[allow(non_snake_case)]
    fn test_zero_ciphertext_proof() {
        let rng = &mut crate::utils::test_rng::get();

        // random ElGamal keypair
        let keypair = SolanaElGamalKeypair::new_rand();
        let mut ciphertext = keypair.pubkey().encrypt(0u32);

        // we flip a coin and generate an invalid proof if the bit is false
        let is_valid_proof = rng.gen_bool(0.5);
        if !is_valid_proof {
            let rand = rng.next_u32() + 1;
            ciphertext = ciphertext.add_amount(rand);
        }

        let solana_proof_data = SolanaZeroCiphertextProofData::new(&keypair, &ciphertext).unwrap();
        assert_eq!(solana_proof_data.verify_proof().is_ok(), is_valid_proof);

        let mut expr_store = IRBuilder::new(true);

        // add inputs
        let mut input_vals = FxHashMap::<usize, EvalValue>::default();
        let _ = <IRBuilder as ExprStore<ScalarField>>::push_curve(
            &mut expr_store,
            CurveExpr::Input(0, Rc::new(InputInfo::from(InputKind::Plaintext))),
        );
        input_vals.insert(
            0,
            EvalValue::Curve(CurvePoint::new(
                <Curve25519Ristretto as AsyncMPCCurve>::Point::from_bytes(
                    &keypair.pubkey().get_point().to_bytes(),
                )
                .unwrap(),
            )),
        );
        let _ = expr_store.push_field(FieldExpr::Input(
            1,
            FieldBounds::<ScalarField>::All.as_input_info(InputKind::Secret),
        ));
        input_vals.insert(
            1,
            EvalValue::Scalar(ScalarField::from(*keypair.secret().get_scalar())),
        );
        let _ = <IRBuilder as ExprStore<ScalarField>>::push_curve(
            &mut expr_store,
            CurveExpr::Input(2, Rc::new(InputInfo::from(InputKind::Plaintext))),
        );
        input_vals.insert(
            2,
            EvalValue::Curve(CurvePoint::new(
                <Curve25519Ristretto as AsyncMPCCurve>::Point::from_bytes(
                    &ciphertext.commitment.get_point().to_bytes(),
                )
                .unwrap(),
            )),
        );
        let _ = <IRBuilder as ExprStore<ScalarField>>::push_curve(
            &mut expr_store,
            CurveExpr::Input(3, Rc::new(InputInfo::from(InputKind::Plaintext))),
        );
        input_vals.insert(
            3,
            EvalValue::Curve(CurvePoint::new(
                <Curve25519Ristretto as AsyncMPCCurve>::Point::from_bytes(
                    &ciphertext.handle.get_point().to_bytes(),
                )
                .unwrap(),
            )),
        );

        let outputs = with_local_expr_store_as_global(
            || {
                let pubkey = CurveValue::new(0);
                let secret = FieldValue::<ScalarField>::from_id(1);
                let commitment = CurveValue::new(2);
                let handle = CurveValue::new(3);

                let arcis_proof_data = ZeroCiphertextProofData::new(
                    &ElGamalKeypair::new_from_inner(
                        ElGamalPubkey::new_from_inner(pubkey),
                        ElGamalSecretKey::new(secret),
                    ),
                    &ElGamalCiphertext {
                        commitment: PedersenCommitment::new(commitment),
                        handle: DecryptHandle::new_from_inner(handle),
                    },
                );

                arcis_proof_data
                    .to_bytes()
                    .into_iter()
                    .map(|byte| FieldValue::<ScalarField>::from(byte).get_id())
                    .collect::<Vec<usize>>()
            },
            &mut expr_store,
        );

        let ir = expr_store.into_ir(outputs);
        let result = ir
            .eval(rng, &mut input_vals)
            .map(|x| {
                x.into_iter()
                    .map(ScalarField::unwrap)
                    .collect::<Vec<ScalarField>>()
            })
            .unwrap();

        let arcis_proof_data_bytes = result
            .iter()
            .map(|byte| byte.to_le_bytes()[0])
            .collect::<Vec<u8>>();

        let arcis_proof_data =
            SolanaZeroCiphertextProofData::from_bytes(&arcis_proof_data_bytes).unwrap();

        let arcis_verification = arcis_proof_data.verify_proof();

        assert_eq!(arcis_verification.is_ok(), is_valid_proof);
    }
}
