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

use crate::{
    core::{
        circuits::boolean::{boolean_value::BooleanValue, byte::Byte},
        global_value::{
            curve_value::{CompressedCurveValue, CurveValue},
            value::FieldValue,
        },
    },
    traits::{GetBit, Invert, Reveal, Select, ToLeBytes},
    utils::{
        field::ScalarField,
        zkp::{
            generators::RangeProofGens,
            inner_product::InnerProductProof,
            pedersen::{Pedersen, PedersenOpening},
            range_proof::batched_range_proof::MAX_SINGLE_BIT_LENGTH,
            transcript::Transcript,
            util::{self, UNIT_LEN},
        },
    },
};
use std::{iter, sync::LazyLock};
use zk_elgamal_proof::encryption::pedersen::H;

pub mod batched_range_proof;
pub mod batched_range_proof_u128;
pub mod batched_range_proof_u256;
pub mod batched_range_proof_u64;

#[allow(non_snake_case, dead_code)]
#[derive(Debug, Clone)]
pub struct RangeProof {
    /// A commitment to the bit-vectors `a_L` and `a_R`.
    pub(crate) A: CompressedCurveValue, // 32 bytes
    /// A commitment to the blinding vectors `s_L` and `s_R`.
    pub(crate) S: CompressedCurveValue, // 32 bytes
    /// A commitment to the `t_1` coefficient of the polynomial `t(x)`.
    pub(crate) T_1: CompressedCurveValue, // 32 bytes
    /// A commitment to the `t_2` coefficient of the polynomial `t(x)`.
    pub(crate) T_2: CompressedCurveValue, // 32 bytes
    /// The evaluation of the polynomial `t(x)` at the challenge point `x`.
    pub(crate) t_x: FieldValue<ScalarField>, // 32 bytes
    /// The blinding factor for the `t_x` value.
    pub(crate) t_x_blinding: FieldValue<ScalarField>, // 32 bytes
    /// The blinding factor for the synthetic commitment to `l(x)` and `r(x)`.
    pub(crate) e_blinding: FieldValue<ScalarField>, // 32 bytes
    /// The inner product proof.
    pub(crate) ipp_proof: InnerProductProof, // 448 bytes for withdraw; 512 for transfer
}

#[allow(non_snake_case)]
impl RangeProof {
    /// Creates an aggregated range proof for a set of values that are committed to a set of
    /// Pedersen commitments.
    ///
    /// This function implements the aggregated proof generation logic from Section 4.3
    /// of the Bulletproofs paper. It allows proving that multiple values are in their
    /// respective ranges, creating one proof that is much smaller than the sum of
    /// individual proofs.
    ///
    /// # Panics
    /// This function will panic if the `openings` vector does not contain the same number
    /// of elements as the `amounts` and `bit_lengths` vectors.
    pub fn new(
        amounts: Vec<FieldValue<ScalarField>>,
        bit_lengths: Vec<usize>,
        openings: Vec<&PedersenOpening>,
        transcript: &mut Transcript<BooleanValue>,
    ) -> Self {
        // 1. Validate inputs
        let m = amounts.len();
        assert!(bit_lengths.len() == m && openings.len() == m);

        // each bit length must be greater than 0 for the proof to make sense
        assert!(bit_lengths
            .iter()
            .all(|bit_length| { *bit_length > 0 && *bit_length <= MAX_SINGLE_BIT_LENGTH }));

        // total vector dimension to compute the ultimate inner product proof for
        let nm: usize = bit_lengths.iter().sum();
        assert!(nm.is_power_of_two());

        let bp_gens = RangeProofGens::new(nm);

        transcript.range_proof_domain_separator(nm as u64);

        // 2. Create commitments A and S.
        // A is a commitment to the bit-vectors a_L and a_R
        let a_blinding = FieldValue::<ScalarField>::random();
        let arcis_H = CurveValue::from(*LazyLock::force(&H));
        let mut A = a_blinding * arcis_H;

        let mut gens_iter = bp_gens.G(nm).zip(bp_gens.H(nm));
        for (amount_i, n_i) in amounts.iter().zip(bit_lengths.iter()) {
            for j in 0..(*n_i) {
                let (G_ij, H_ij) = gens_iter.next().unwrap();
                let v_ij = amount_i.get_bit(j, false);
                let mut point = -*H_ij;
                // Add G_ij if bit is 1, else do nothing (since a_R = a_L - 1)
                point = v_ij.select(*G_ij, point);
                A += point;
            }
        }
        let A = A.reveal().compress();

        // generate blinding factors and generate their Pedersen vector commitment
        let s_L = (0..nm)
            .map(|_| FieldValue::<ScalarField>::random())
            .collect::<Vec<FieldValue<ScalarField>>>();
        let s_R = (0..nm)
            .map(|_| FieldValue::<ScalarField>::random())
            .collect::<Vec<FieldValue<ScalarField>>>();

        // generate blinding factor for Pedersen commitment; `s_blinding` should not to be confused
        // with blinding factors for the actual inner product vector
        let s_blinding = FieldValue::<ScalarField>::random();

        let S = CurveValue::multiscalar_mul(
            iter::once(&s_blinding)
                .chain(s_L.iter())
                .chain(s_R.iter())
                .copied()
                .collect::<Vec<FieldValue<ScalarField>>>(),
            iter::once(&arcis_H)
                .chain(bp_gens.G(nm))
                .chain(bp_gens.H(nm))
                .copied()
                .collect::<Vec<CurveValue>>(),
        )
        .reveal()
        .compress();

        // 3. Derive challenges y and z.
        transcript.append_point(b"A", &A);
        transcript.append_point(b"S", &S);

        let y = transcript.challenge_scalar(b"y");
        let z = transcript.challenge_scalar(b"z");

        // 4. Construct the blinded vector polynomials l(x) and r(x). l(x) = (a_L - z*1) + s_L*x
        //    r(x) = y^nm o (a_R + z*1 + s_R*x) + (z^2*2^n_1 || ... || z^{m+1}*2^n_m) where `o` is
        //    the Hadamard product and `||` is vector concatenation.
        let mut l_poly = util::VecPoly1::zero(nm);
        let mut r_poly = util::VecPoly1::zero(nm);

        let mut i = 0;
        let mut exp_z = z * z;
        let mut exp_y = FieldValue::<ScalarField>::from(1);

        for (amount_i, n_i) in amounts.iter().zip(bit_lengths.iter()) {
            let mut exp_2 = FieldValue::<ScalarField>::from(1);

            for j in 0..(*n_i) {
                let a_L_j = FieldValue::<ScalarField>::from(amount_i.get_bit(j, false));
                let a_R_j = a_L_j - FieldValue::<ScalarField>::from(1);

                l_poly.0[i] = a_L_j - z;
                l_poly.1[i] = s_L[i];
                r_poly.0[i] = exp_y * (a_R_j + z) + exp_z * exp_2;
                r_poly.1[i] = exp_y * s_R[i];

                exp_y *= y;
                exp_2 = exp_2 + exp_2;

                // `i` is capped by the sum of vectors in `bit_lengths`
                i = i.checked_add(1).unwrap();
            }
            exp_z *= z;
        }

        // 5. Compute the inner product polynomial t(x) = <l(x), r(x)>.
        let t_poly = l_poly.inner_product(&r_poly).unwrap();

        // 6. Commit to the t_1 and t_2 coefficients of t(x).
        let (T_1, t_1_blinding) = Pedersen::new(t_poly.1);
        let (T_2, t_2_blinding) = Pedersen::new(t_poly.2);

        let T_1 = T_1.get_point().reveal().compress();
        let T_2 = T_2.get_point().reveal().compress();

        transcript.append_point(b"T_1", &T_1);
        transcript.append_point(b"T_2", &T_2);

        // 7. Derive challenge x and compute openings.
        let x = transcript.challenge_scalar(b"x");

        // Compute the aggregated blinding factor for all value commitments.
        let mut agg_opening = FieldValue::<ScalarField>::from(0);
        let mut exp_z = z;
        for opening in openings {
            exp_z *= z;
            agg_opening += exp_z * *opening.get_scalar();
        }

        let t_blinding_poly = util::Poly2(
            agg_opening,
            *t_1_blinding.get_scalar(),
            *t_2_blinding.get_scalar(),
        );

        let t_x = t_poly.eval(x).reveal();
        let t_x_blinding = t_blinding_poly.eval(x).reveal();

        transcript.append_scalar(b"t_x", &t_x);
        transcript.append_scalar(b"t_x_blinding", &t_x_blinding);

        // homomorphically compuate the openings for A + x*S
        let e_blinding = (a_blinding + s_blinding * x).reveal();

        // 8. Finally, create the inner product proof.
        let l_vec = l_poly.eval(x);
        let r_vec = r_poly.eval(x);

        transcript.append_scalar(b"e_blinding", &e_blinding);

        // compute the inner product argument on the commitment:
        // P = <l(x), G> + <r(x), H'> + <l(x), r(x)>*Q
        let w = transcript.challenge_scalar(b"w");
        let Q = w * CurveValue::generator();

        let G_factors = iter::repeat_n(FieldValue::<ScalarField>::from(1), nm)
            .collect::<Vec<FieldValue<ScalarField>>>();
        // on plaintext values we simply set is_expected_non_zero = true
        let H_factors = util::exp_iter(y.invert(true))
            .take(nm)
            .collect::<Vec<FieldValue<ScalarField>>>();

        // this variable exists for backwards compatibility
        let _c = transcript.challenge_scalar(b"c");

        let ipp_proof = InnerProductProof::new(
            &Q,
            &G_factors,
            &H_factors,
            bp_gens.G(nm).cloned().collect(),
            bp_gens.H(nm).cloned().collect(),
            l_vec,
            r_vec,
            transcript,
        );

        // compute challenge `d` for consistency with the verifier
        transcript.append_scalar(b"ipp_a", &ipp_proof.a);
        transcript.append_scalar(b"ipp_b", &ipp_proof.b);
        let _d = transcript.challenge_scalar(b"d");

        RangeProof {
            A,
            S,
            T_1,
            T_2,
            t_x,
            t_x_blinding,
            e_blinding,
            ipp_proof,
        }
    }

    pub fn to_bytes(&self) -> Vec<Byte<BooleanValue>> {
        let mut buf = Vec::with_capacity(7 * UNIT_LEN + self.ipp_proof.serialized_size());
        buf.extend_from_slice(&self.A.to_bytes());
        buf.extend_from_slice(&self.S.to_bytes());
        buf.extend_from_slice(&self.T_1.to_bytes());
        buf.extend_from_slice(&self.T_2.to_bytes());
        buf.extend_from_slice(&self.t_x.to_le_bytes());
        buf.extend_from_slice(&self.t_x_blinding.to_le_bytes());
        buf.extend_from_slice(&self.e_blinding.to_le_bytes());
        buf.extend_from_slice(&self.ipp_proof.to_bytes());
        buf
    }
}
