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

use crate::{
    core::{
        circuits::boolean::{boolean_value::BooleanValue, byte::Byte},
        global_value::{curve_value::CurveValue, value::FieldValue},
    },
    utils::{
        field::ScalarField,
        zkp::{
            pedersen::{PedersenCommitment, PedersenOpening},
            transcript::Transcript,
        },
    },
};
use zk_elgamal_proof::encryption::PEDERSEN_COMMITMENT_LEN;

// same as zk_elgamal_proof_program::proof_data::batched_range_proof::MAX_COMMITMENTS,
// but currently not public
const MAX_COMMITMENTS: usize = 8;

pub const BATCHED_RANGE_PROOF_CONTEXT_LEN: usize =
    MAX_COMMITMENTS * PEDERSEN_COMMITMENT_LEN + MAX_COMMITMENTS;

/// A bit length in a batched range proof must be at most 64.
pub const MAX_SINGLE_BIT_LENGTH: usize = 64;

/// The context data needed to verify a range-proof for a Pedersen committed value.
///
/// The context data is shared by all `VerifyBatchedRangeProof{N}` instructions.
#[derive(Clone, Copy)]
pub struct BatchedRangeProofContext {
    pub commitments: [PedersenCommitment; MAX_COMMITMENTS],
    pub bit_lengths: [u8; MAX_COMMITMENTS],
}

impl BatchedRangeProofContext {
    pub(crate) fn new_transcript(&self) -> Transcript<BooleanValue> {
        let mut transcript = Transcript::new(b"batched-range-proof-instruction");
        transcript.append_message(
            b"commitments",
            &self
                .commitments
                .iter()
                .flat_map(|commitment| commitment.get_point().compress().to_bytes().to_vec())
                .collect::<Vec<Byte<BooleanValue>>>(),
        );
        transcript.append_message(
            b"bit-lengths",
            &self
                .bit_lengths
                .iter()
                .copied()
                .map(Byte::from)
                .collect::<Vec<Byte<BooleanValue>>>(),
        );
        transcript
    }

    // TODO: the nodes must check that none of commitments is the identity
    pub(crate) fn new(
        commitments: &[&PedersenCommitment],
        amounts: &[FieldValue<ScalarField>],
        bit_lengths: &[usize],
        openings: &[&PedersenOpening],
    ) -> Self {
        // the number of commitments is capped at 8
        let num_commitments = commitments.len();
        assert!(
            num_commitments <= MAX_COMMITMENTS
                && num_commitments == amounts.len()
                && num_commitments == bit_lengths.len()
                && num_commitments == openings.len()
        );

        let mut all_commitments = [PedersenCommitment::new(CurveValue::default()); MAX_COMMITMENTS];
        all_commitments[0..commitments.len()].copy_from_slice(
            &commitments
                .iter()
                .map(|&commitment| *commitment)
                .collect::<Vec<PedersenCommitment>>(),
        );

        let mut all_bit_lengths = [0; MAX_COMMITMENTS];
        all_bit_lengths[..bit_lengths.len()]
            .copy_from_slice(&bit_lengths.iter().map(|&b| b as u8).collect::<Vec<u8>>());

        BatchedRangeProofContext {
            commitments: all_commitments,
            bit_lengths: all_bit_lengths,
        }
    }

    pub fn to_bytes(&self) -> [Byte<BooleanValue>; BATCHED_RANGE_PROOF_CONTEXT_LEN] {
        let mut bytes = [Byte::<BooleanValue>::from(0); BATCHED_RANGE_PROOF_CONTEXT_LEN];
        bytes[..MAX_COMMITMENTS * PEDERSEN_COMMITMENT_LEN].copy_from_slice(
            &self
                .commitments
                .iter()
                .flat_map(|commitment| commitment.get_point().compress().to_bytes().to_vec())
                .collect::<Vec<Byte<BooleanValue>>>(),
        );
        bytes[MAX_COMMITMENTS * PEDERSEN_COMMITMENT_LEN..].copy_from_slice(
            &self
                .bit_lengths
                .iter()
                .copied()
                .map(Byte::from)
                .collect::<Vec<Byte<BooleanValue>>>(),
        );
        bytes
    }
}
