//! Zero-knowledge-aware wrapper for [`BasefoldVerifier`] with custom batch evaluation support.
//!
//! This module provides [`ZkStackedPcsVerifier`], a wrapper around `StackedPcsVerifier` that
//! verifies Basefold proofs for pre-batched, extension-field MLEs. The code is largely
//! copy-pasted from the standard Basefold verifier with two key modifications:
//!
//! 1. **Extension field openings**: The commitment is over a 4-column base-field tensor
//!    (representing the base-field components of extension-field coefficients), so the
//!    FRI query phase works with extension-field evaluations reconstructed from base-field
//!    openings.
//!
//! 2. **Closure-based batch evaluation**: Instead of computing batch evaluations internally,
//!    `verify_trusted_ext_mle_evaluation` accepts a `compute_batch_evals` closure. This
//!    allows the caller (e.g., `verifier.rs`) to inject protocol-specific correction logic
//!    (such as RLC padding correction) without polluting this copy-pasted Basefold code.
//!
//! # Relationship to Other Modules
//!
//! - **Prover side**: Proofs verified here are generated by
//!   [`prove_from_batched_inputs`](crate::zk::stacked_pcs::basefold_prover_wrapper::prove_from_batched_inputs)
//!   in the prover wrapper.
//! - **Caller**: [`verifier.rs`](crate::zk::stacked_pcs::verifier) constructs the `compute_batch_evals` closure
//!   that captures the RLC padding vector and expected evaluations, then passes it to
//!   `verify_trusted_ext_mle_evaluation`.
//!
//! # Usage
//!
//! ```ignore
//! let verifier = ZkStackedPcsVerifier::new(stacked_verifier);
//!
//! // The caller constructs a closure that computes corrected batch evaluations
//! // from query indices and the tensor height.
//! let compute_batch_evals = |query_indices: &[usize], log_tensor_height: usize| -> Vec<GC::EF> {
//!     // ... correction logic using RLC padding, expected evals, etc.
//! };
//!
//! verifier.verify_trusted_ext_mle_evaluation(
//!     &commitment,
//!     eval_point,
//!     eval_claim,
//!     &proof,
//!     compute_batch_evals,
//!     &mut challenger,
//! )?;
//! ```

use crate::zk::inner::{ZkIopCtx, ZkMerkleizer};
use itertools::Itertools;
use slop_algebra::{AbstractExtensionField, AbstractField, TwoAdicField};
use slop_basefold::{BaseFoldVerifierError, BasefoldProof, BasefoldVerifier, FriConfig};
use slop_basefold_prover::BasefoldProver;
use slop_challenger::{CanObserve, CanSampleBits, FieldChallenger, GrindingChallenger, IopCtx};
use slop_merkle_tree::MerkleTreeTcsError;
use slop_multilinear::Point;
use slop_stacked::StackedPcsVerifier;
use slop_utils::reverse_bits_len;

use super::basefold_prover_wrapper::ZkBasefoldProver;

/// Zero-knowledge-aware wrapper around [`StackedPcsVerifier`].
///
/// This struct wraps the standard stacked PCS verifier and provides a modified
/// Basefold verification routine that:
/// - Operates on extension-field MLE commitments (stored as base-field tensors)
/// - Accepts a closure for computing corrected batch evaluations during FRI queries
///
/// The wrapper only adds `verify_trusted_ext_mle_evaluation` and `verify_queries`;
/// all other verification logic lives in the caller (`verifier.rs`).
pub struct ZkStackedPcsVerifier<GC>
where
    GC: IopCtx<F: TwoAdicField>,
{
    /// The underlying stacked PCS verifier (contains `BasefoldVerifier` and stacking config)
    pub inner: StackedPcsVerifier<GC>,
}

impl<GC: IopCtx> ZkStackedPcsVerifier<GC>
where
    GC::F: TwoAdicField,
    GC::EF: TwoAdicField,
{
    pub const fn new(inner: StackedPcsVerifier<GC>) -> Self {
        Self { inner }
    }

    /// Verifies a Basefold proof for a pre-batched, extension-field MLE evaluation.
    ///
    /// This is the counterpart to
    /// [`prove_from_batched_inputs`](crate::basefold_prover_wrapper::prove_from_batched_inputs)
    /// on the prover side. The commitment is over a base-field tensor whose columns are the
    /// base-field components of the extension-field MLE coefficients.
    ///
    /// # Parameters
    /// - `commitment`: Merkle root of the committed (base-field) tensor
    /// - `point`: Evaluation point in the extension field
    /// - `eval_claim`: Claimed evaluation of the batched MLE at `point`
    /// - `proof`: The Basefold proof to verify
    /// - `compute_batch_evals`: Closure called with `(query_indices, log_tensor_height)` that
    ///   returns corrected batch evaluations for the FRI query phase. This is where the caller
    ///   injects protocol-specific logic (e.g., subtracting RLC padding contributions).
    /// - `challenger`: Fiat-Shamir challenger (state must match the prover's)
    ///
    /// # Verification Steps
    /// 1. Shape checks on the proof structure
    /// 2. Sumcheck verification: checks round-by-round consistency of univariate messages
    /// 3. Proof-of-work verification (grinding)
    /// 4. FRI query phase: samples query indices, calls `compute_batch_evals` to get
    ///    corrected evaluations, verifies Merkle openings, and checks FRI consistency
    /// 5. Final polynomial consistency check
    pub fn verify_trusted_ext_mle_evaluation(
        &self,
        commitments: &[GC::Digest],
        mut point: Point<GC::EF>,
        eval_claim: GC::EF,
        proof: &BasefoldProof<GC>,
        compute_batch_evals: impl FnOnce(&[usize], usize) -> Vec<GC::EF>,
        challenger: &mut GC::Challenger,
    ) -> Result<(), BaseFoldVerifierError<MerkleTreeTcsError>> {
        let verifier = &self.inner.basefold_verifier;

        if proof.component_polynomials_query_openings_and_proofs.len() != commitments.len() {
            return Err(BaseFoldVerifierError::IncorrectShape);
        }

        // Assert correctness of shape.
        if proof.fri_commitments.len() != proof.univariate_messages.len()
            || proof.fri_commitments.len() != point.dimension()
            || proof.univariate_messages.is_empty()
        {
            return Err(BaseFoldVerifierError::SumcheckFriLengthMismatch);
        }

        // The prover messages correspond to fixing the last coordinate first, so we reverse the
        // underlying point for the verification.
        point.reverse();

        // Sample the challenges used for FRI folding and BaseFold random linear combinations.
        let betas = proof
            .fri_commitments
            .iter()
            .zip_eq(proof.univariate_messages.iter())
            .map(|(commitment, poly)| {
                poly.iter().copied().for_each(|x| challenger.observe_ext_element(x));
                challenger.observe(*commitment);
                challenger.sample_ext_element::<GC::EF>()
            })
            .collect::<Vec<_>>();

        // Check the consistency of the first univariate message with the claimed evaluation. The
        // first_poly is supposed to be `vals(X_0, X_1, ..., X_{d-1}, 0), vals(X_0, X_1, ...,
        // X_{d-1}, 1)`. Given this, the claimed evaluation should be `(1 - X_d) *
        // first_poly[0] + X_d * first_poly[1]`.
        let first_poly = proof.univariate_messages[0];
        if eval_claim != (GC::EF::one() - *point[0]) * first_poly[0] + *point[0] * first_poly[1] {
            return Err(BaseFoldVerifierError::Sumcheck);
        };

        // Fold the two messages into a single evaluation claim for the next round, using the
        // sampled randomness.
        let mut expected_eval = first_poly[0] + betas[0] * first_poly[1];

        // Check round-by-round consistency between the successive sumcheck univariate messages.
        for (i, (poly, beta)) in
            proof.univariate_messages[1..].iter().zip_eq(betas[1..].iter()).enumerate()
        {
            // The check is similar to the one for `first_poly`.
            let i = i + 1;
            if expected_eval != (GC::EF::one() - *point[i]) * poly[0] + *point[i] * poly[1] {
                return Err(BaseFoldVerifierError::Sumcheck);
            }

            // Fold the two pieces of the message.
            expected_eval = poly[0] + *beta * poly[1];
        }

        challenger.observe_ext_element(proof.final_poly);

        // Check proof of work (grinding to find a number that hashes to have
        // `verifier.fri_config.proof_of_work_bits` zeroes at the beginning).
        if !challenger.check_witness(verifier.fri_config.proof_of_work_bits, proof.pow_witness) {
            return Err(BaseFoldVerifierError::Pow);
        }

        let log_len = proof.fri_commitments.len();

        if log_len + verifier.fri_config.log_blowup() > GC::F::TWO_ADICITY {
            return Err(BaseFoldVerifierError::TwoAdicityOverflow);
        }

        // Sample query indices for the FRI query IOPP part of BaseFold. This part is very similar
        // to the corresponding part in the univariate FRI verifier.
        let query_indices = (0..verifier.fri_config.num_queries)
            .map(|_| challenger.sample_bits(log_len + verifier.fri_config.log_blowup()))
            .collect::<Vec<_>>();

        // Shape-check and verify Merkle openings for each committed polynomial.
        let log_tensor_height = log_len + verifier.fri_config.log_blowup();
        for (opening_and_proof, commitment) in
            proof.component_polynomials_query_openings_and_proofs.iter().zip_eq(commitments)
        {
            let initial_opening_values = &opening_and_proof.values;
            if initial_opening_values.dimensions.sizes().len() != 2 {
                return Err(BaseFoldVerifierError::IncorrectShape);
            }
            if initial_opening_values.dimensions.sizes()[0] != query_indices.len() {
                return Err(BaseFoldVerifierError::IncorrectShape);
            }
            if opening_and_proof.proof.log_tensor_height != log_tensor_height {
                return Err(BaseFoldVerifierError::IncorrectShape);
            }

            verifier
                .tcs
                .verify_tensor_openings(
                    commitment,
                    &query_indices,
                    &opening_and_proof.values,
                    initial_opening_values.sizes()[1],
                    log_tensor_height,
                    &opening_and_proof.proof,
                )
                .map_err(BaseFoldVerifierError::TcsError)?;
        }

        // Compute corrected batch evaluations via the caller-provided closure.
        let batch_evals_corrected = compute_batch_evals(&query_indices, log_tensor_height);

        // Check that the query openings are consistent as FRI messages.
        self.verify_queries(
            &proof.fri_commitments,
            &query_indices,
            proof.final_poly,
            batch_evals_corrected,
            &proof.query_phase_openings_and_proofs,
            &betas,
        )?;

        if proof.final_poly
            != proof.univariate_messages.last().unwrap()[0]
                + *betas.last().unwrap() * proof.univariate_messages.last().unwrap()[1]
        {
            return Err(BaseFoldVerifierError::SumcheckFinalPolyMismatch);
        }

        Ok(())
    }

    /// Verifies the FRI query phase for all sampled query indices.
    ///
    /// Starting from the `reduced_openings` (corrected batch evaluations at the leaf level),
    /// this function iteratively folds evaluations through each FRI round, checking that:
    /// - Each folded evaluation is consistent with the Merkle-opened sibling value
    /// - Each round's Merkle opening verifies against the corresponding commitment
    /// - All queries converge to the declared `final_poly` constant
    ///
    /// This is adapted from the standard Plonky3 FRI verifier to work with a single
    /// committed vector (rather than multiple).
    fn verify_queries(
        &self,
        commitments: &[GC::Digest],
        indices: &[usize],
        final_poly: GC::EF,
        reduced_openings: Vec<GC::EF>,
        query_openings: &[slop_merkle_tree::MerkleTreeOpeningAndProof<GC>],
        betas: &[GC::EF],
    ) -> Result<(), BaseFoldVerifierError<MerkleTreeTcsError>> {
        let verifier = &self.inner.basefold_verifier;
        let log_max_height = commitments.len() + verifier.fri_config.log_blowup();

        let mut folded_evals = reduced_openings;
        let mut indices = indices.to_vec();

        let mut xis = indices
            .iter()
            .map(|index| {
                GC::F::two_adic_generator(log_max_height)
                    .exp_u64(reverse_bits_len(*index, log_max_height) as u64)
            })
            .collect::<Vec<_>>();

        if commitments.len() != query_openings.len() || commitments.len() != betas.len() {
            return Err(BaseFoldVerifierError::IncorrectShape);
        }

        // Loop over the FRI queries.
        for (round_idx, ((commitment, query_opening), beta)) in (verifier.fri_config.log_blowup()
            ..log_max_height)
            .rev()
            .zip_eq(commitments.iter().zip_eq(query_openings.iter()).zip_eq(betas))
        {
            let openings = &query_opening.values;
            if openings.dimensions.sizes().len() != 2 {
                return Err(BaseFoldVerifierError::IncorrectShape);
            }

            if indices.len() != folded_evals.len()
                || indices.len() != openings.dimensions.sizes()[0]
                || indices.len() != xis.len()
            {
                return Err(BaseFoldVerifierError::IncorrectShape);
            }

            for (((index, folded_eval), opening), x) in indices
                .iter_mut()
                .zip_eq(folded_evals.iter_mut())
                .zip_eq(openings.split())
                .zip_eq(xis.iter_mut())
            {
                let index_sibling = *index ^ 1;
                let index_pair = *index >> 1;

                if opening.total_len() != 2 * <GC::EF as AbstractExtensionField<GC::F>>::D {
                    return Err(BaseFoldVerifierError::IncorrectShape);
                }

                let evals: [GC::EF; 2] = opening
                    .as_slice()
                    .chunks_exact(GC::EF::D)
                    .map(GC::EF::from_base_slice)
                    .collect::<Vec<_>>()
                    .try_into()
                    .unwrap();

                // Check that the folded evaluation is consistent with the FRI query proof opening.
                if evals[*index % 2] != *folded_eval {
                    return Err(BaseFoldVerifierError::QueryValueMismatch);
                }

                let mut xs = [*x; 2];
                xs[index_sibling % 2] *= GC::F::two_adic_generator(1);

                // interpolate and evaluate at beta
                *folded_eval = evals[0]
                    + (*beta - xs[0].into()) * (evals[1] - evals[0]) / GC::EF::from(xs[1] - xs[0]);

                *index = index_pair;
                *x = x.square();
            }

            // The magic constant 2 here is the folding factor we use for FRI.
            if round_idx != query_opening.proof.log_tensor_height
                || query_opening.proof.width != GC::EF::D * 2
            {
                return Err(BaseFoldVerifierError::IncorrectShape);
            }

            // Check that the opening is consistent with the commitment.
            verifier
                .tcs
                .verify_tensor_openings(
                    commitment,
                    &indices,
                    &query_opening.values,
                    GC::EF::D * 2,
                    round_idx,
                    &query_opening.proof,
                )
                .map_err(BaseFoldVerifierError::TcsError)?;
        }

        for folded_eval in folded_evals {
            if folded_eval != final_poly {
                return Err(BaseFoldVerifierError::QueryFinalPolyMismatch);
            }
        }

        Ok(())
    }
}

/// Creates both a `ZkBasefoldProver` and `ZkStackedPcsVerifier` with default FRI configuration.
///
/// This is a convenience function that simplifies the typical setup where you need both
/// a prover and verifier. The underlying `BasefoldVerifier` is shared (prover borrows it
/// before verifier takes ownership).
///
/// The `num_encoding_variables` value is fixed here and all subsequent
/// [`commit_mle`](crate::zk::ZkProverCtx::commit_mle) and
/// [`read_oracle`](crate::compiler::ReadingCtx::read_oracle) calls must use a matching
/// `num_encoding_variables` (inferred from the MLE size for `commit_mle`, or passed
/// directly for `read_oracle`).
///
/// # Arguments
/// * `num_expected_commitments` — upper bound on the number of MLE commitments that
///   will be made during the protocol.
/// * `num_encoding_variables` — number of variables per stacked polynomial (encoding
///   width). Each committed MLE will be stacked into a tensor whose rows have
///   `2^num_encoding_variables` entries.
///
/// # Example
/// ```ignore
/// let (pcs_prover, pcs_verifier) =
///     initialize_zk_prover_and_verifier::<GC, MK>(1, NUM_ENCODING_VARIABLES);
/// ```
pub fn initialize_zk_prover_and_verifier<GC: ZkIopCtx, MK: ZkMerkleizer<GC>>(
    num_expected_commitments: usize,
    num_encoding_variables: u32,
) -> (ZkBasefoldProver<GC, MK>, ZkStackedPcsVerifier<GC>)
where
    GC::F: TwoAdicField,
{
    let fri_config = FriConfig::default_fri_config();
    let basefold_verifier = BasefoldVerifier::<GC>::new(fri_config, num_expected_commitments);
    let basefold_prover = BasefoldProver::new(&basefold_verifier);
    let zk_basefold_prover = ZkBasefoldProver::new(basefold_prover);
    let stacked_verifier = StackedPcsVerifier::new(basefold_verifier, num_encoding_variables);
    let zk_stacked_verifier = ZkStackedPcsVerifier::new(stacked_verifier);
    (zk_basefold_prover, zk_stacked_verifier)
}
