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

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

pub const CIPHERTEXT_COMMITMENT_EQUALITY_PROOF_LEN: usize = 192;

pub const CIPHERTEXT_COMMITMENT_EQUALITY_PROOF_CONTEXT_LEN: usize =
    ELGAMAL_PUBKEY_LEN + ELGAMAL_CIPHERTEXT_LEN + PEDERSEN_COMMITMENT_LEN;

pub const CIPHERTEXT_COMMITMENT_EQUALITY_PROOF_DATA_LEN: usize =
    CIPHERTEXT_COMMITMENT_EQUALITY_PROOF_CONTEXT_LEN + CIPHERTEXT_COMMITMENT_EQUALITY_PROOF_LEN;

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

/// The context data needed to verify a ciphertext-commitment equality proof.
#[derive(Clone, Copy)]
pub struct CiphertextCommitmentEqualityProofContext {
    /// The ElGamal pubkey
    pub pubkey: ElGamalPubkey,
    /// The ciphertext encrypted under the ElGamal pubkey
    pub ciphertext: ElGamalCiphertext,
    /// The Pedersen commitment
    pub commitment: PedersenCommitment,
}

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

    pub fn to_bytes(
        &self,
    ) -> [Byte<BooleanValue>; CIPHERTEXT_COMMITMENT_EQUALITY_PROOF_CONTEXT_LEN] {
        let mut bytes =
            [Byte::<BooleanValue>::from(0); CIPHERTEXT_COMMITMENT_EQUALITY_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());
        offset += DECRYPT_HANDLE_LEN;
        bytes[offset..offset + PEDERSEN_COMMITMENT_LEN]
            .copy_from_slice(&self.commitment.get_point().compress().to_bytes());
        bytes
    }
}

impl CiphertextCommitmentEqualityProofData {
    pub fn new(
        keypair: &ElGamalKeypair,
        ciphertext: &ElGamalCiphertext,
        commitment: &PedersenCommitment,
        opening: &PedersenOpening,
        amount: FieldValue<ScalarField>,
    ) -> Self {
        let context = CiphertextCommitmentEqualityProofContext {
            pubkey: *keypair.pubkey(),
            ciphertext: *ciphertext,
            commitment: *commitment,
        };
        let mut transcript = context.new_transcript();
        let proof = CiphertextCommitmentEqualityProof::new(
            keypair,
            ciphertext,
            opening,
            amount,
            &mut transcript,
        );
        CiphertextCommitmentEqualityProofData { context, proof }
    }

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

/// Equality 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 CiphertextCommitmentEqualityProof {
    Y_0: CompressedCurveValue,
    Y_1: CompressedCurveValue,
    Y_2: CompressedCurveValue,
    z_s: FieldValue<ScalarField>,
    z_x: FieldValue<ScalarField>,
    z_r: FieldValue<ScalarField>,
}

#[allow(non_snake_case)]
impl CiphertextCommitmentEqualityProof {
    /// Creates a ciphertext-commitment equality proof.
    ///
    /// The function does *not* hash the public key, ciphertext, or commitment 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.
    ///
    /// Note that the proof constructor does not take the actual Pedersen commitment as input; it
    /// takes the associated Pedersen opening instead.
    ///
    /// * `keypair` - The ElGamal keypair associated with the first to be proved
    /// * `ciphertext` - The main ElGamal ciphertext to be proved
    /// * `opening` - The opening associated with the main Pedersen commitment to be proved
    /// * `amount` - The message associated with the ElGamal ciphertext and Pedersen commitment
    /// * `transcript` - The transcript that does the bookkeeping for the Fiat-Shamir heuristic
    pub fn new(
        keypair: &ElGamalKeypair,
        ciphertext: &ElGamalCiphertext,
        opening: &PedersenOpening,
        amount: FieldValue<ScalarField>,
        transcript: &mut Transcript<BooleanValue>,
    ) -> Self {
        transcript.ciphertext_commitment_equality_proof_domain_separator();

        // extract the relevant scalar and Ristretto points from the inputs
        let P = keypair.pubkey().get_point();
        let D = ciphertext.handle.get_point();

        let s = keypair.secret().get_scalar();
        let x = amount;
        let r = opening.get_scalar();

        // generate random masking factors that also serves as nonces
        let y_s = FieldValue::<ScalarField>::random();
        let y_x = FieldValue::<ScalarField>::random();
        let y_r = FieldValue::<ScalarField>::random();

        let Y_0 = (y_s * *P).reveal().compress();
        let Y_1 = CurveValue::multiscalar_mul(vec![y_x, y_s], vec![CurveValue::generator(), *D])
            .reveal()
            .compress();
        let Y_2 = CurveValue::multiscalar_mul(
            vec![y_x, y_r],
            vec![
                CurveValue::generator(),
                CurveValue::from(*LazyLock::force(&H)),
            ],
        )
        .reveal()
        .compress();

        // record masking factors in the transcript
        transcript.append_point(b"Y_0", &Y_0);
        transcript.append_point(b"Y_1", &Y_1);
        transcript.append_point(b"Y_2", &Y_2);

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

        // compute the masked values
        let z_s = ((c * *s) + y_s).reveal();
        let z_x = ((c * x) + y_x).reveal();
        let z_r = ((c * *r) + y_r).reveal();

        // compute challenge `w` for consistency with verification
        transcript.append_scalar(b"z_s", &z_s);
        transcript.append_scalar(b"z_x", &z_x);
        transcript.append_scalar(b"z_r", &z_r);
        let _w = transcript.challenge_scalar(b"w");

        CiphertextCommitmentEqualityProof {
            Y_0,
            Y_1,
            Y_2,
            z_s,
            z_x,
            z_r,
        }
    }

    pub fn to_bytes(&self) -> [Byte<BooleanValue>; CIPHERTEXT_COMMITMENT_EQUALITY_PROOF_LEN] {
        let mut buf = [Byte::<BooleanValue>::from(0); CIPHERTEXT_COMMITMENT_EQUALITY_PROOF_LEN];
        let mut chunks = buf.chunks_mut(UNIT_LEN);
        chunks.next().unwrap().copy_from_slice(&self.Y_0.to_bytes());
        chunks.next().unwrap().copy_from_slice(&self.Y_1.to_bytes());
        chunks.next().unwrap().copy_from_slice(&self.Y_2.to_bytes());
        chunks
            .next()
            .unwrap()
            .copy_from_slice(&self.z_s.to_le_bytes());
        chunks
            .next()
            .unwrap()
            .copy_from_slice(&self.z_x.to_le_bytes());
        chunks
            .next()
            .unwrap()
            .copy_from_slice(&self.z_r.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,
            used_field::UsedField,
            zkp::{
                ciphertext_commitment_equality::CiphertextCommitmentEqualityProofData,
                elgamal::{
                    DecryptHandle,
                    ElGamalCiphertext,
                    ElGamalKeypair,
                    ElGamalPubkey,
                    ElGamalSecretKey,
                },
                pedersen::{PedersenCommitment, PedersenOpening},
            },
        },
    };
    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,
            pedersen::Pedersen as SolanaPedersen,
        },
        zk_elgamal_proof_program::proof_data::{
            CiphertextCommitmentEqualityProofData as SolanaCiphertextCommitmentEqualityProofData,
            ZkProofData,
        },
    };

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

        // random ElGamal keypair and amount
        let keypair = SolanaElGamalKeypair::new_rand();
        let mut amount = rng.next_u64();
        let ciphertext = keypair.pubkey().encrypt(amount);
        let (mut commitment, mut opening) = SolanaPedersen::new(amount);

        // 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 {
            amount = rng.next_u64();
            (commitment, opening) = SolanaPedersen::new(amount);
        }

        let solana_proof_data = SolanaCiphertextCommitmentEqualityProofData::new(
            &keypair,
            &ciphertext,
            &commitment,
            &opening,
            amount,
        )
        .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 _ = <IRBuilder as ExprStore<ScalarField>>::push_curve(
            &mut expr_store,
            CurveExpr::Input(4, Rc::new(InputInfo::from(InputKind::Plaintext))),
        );
        input_vals.insert(
            4,
            EvalValue::Curve(CurvePoint::new(
                <Curve25519Ristretto as AsyncMPCCurve>::Point::from_bytes(
                    &commitment.get_point().to_bytes(),
                )
                .unwrap(),
            )),
        );
        let _ = expr_store.push_field(FieldExpr::Input(
            5,
            FieldBounds::<ScalarField>::All.as_input_info(InputKind::Secret),
        ));
        input_vals.insert(
            5,
            EvalValue::Scalar(ScalarField::from(*opening.get_scalar())),
        );
        let _ = expr_store.push_field(FieldExpr::Input(
            6,
            FieldBounds::new(
                ScalarField::from(0),
                ScalarField::power_of_two(64) - ScalarField::from(1),
            )
            .as_input_info(InputKind::Secret),
        ));
        input_vals.insert(6, EvalValue::Scalar(ScalarField::from(amount)));

        let outputs = with_local_expr_store_as_global(
            || {
                let pubkey = CurveValue::new(0);
                let secret = FieldValue::<ScalarField>::from_id(1);
                let ciphertext_commitment = CurveValue::new(2);
                let ciphertext_handle = CurveValue::new(3);
                let commitment = CurveValue::new(4);
                let opening = FieldValue::<ScalarField>::from_id(5);
                let amount = FieldValue::<ScalarField>::from_id(6);

                let arcis_proof_data = CiphertextCommitmentEqualityProofData::new(
                    &ElGamalKeypair::new_from_inner(
                        ElGamalPubkey::new_from_inner(pubkey),
                        ElGamalSecretKey::new(secret),
                    ),
                    &ElGamalCiphertext {
                        commitment: PedersenCommitment::new(ciphertext_commitment),
                        handle: DecryptHandle::new_from_inner(ciphertext_handle),
                    },
                    &PedersenCommitment::new(commitment),
                    &PedersenOpening::new(opening),
                    amount,
                );

                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 =
            SolanaCiphertextCommitmentEqualityProofData::from_bytes(&arcis_proof_data_bytes)
                .unwrap();

        let arcis_verification = arcis_proof_data.verify_proof();

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