//! Proof accumulator for batched verification.
//!
//! This module implements the accumulator scheme that allows
//! combining multiple proofs into a single verification.

use std::marker::PhantomData;

use crate::crypto::transcript::Transcript;
use crate::traits::PairingEngine;

/// An accumulated proof that combines multiple individual proofs.
#[derive(Debug, Clone)]
pub struct AccumulatedProof<E: PairingEngine> {
    /// Combined adjusted commitment: Σ r_i * (C_i - [v_i]_1 + z_i * π_i)
    pub adjusted_commitment: E::G1Affine,

    /// Combined quotient commitment: Σ r_i * π_i
    pub combined_quotient: E::G1Affine,

    /// Number of proofs accumulated.
    pub count: usize,
}

/// Instance for accumulation (a full KZG opening).
#[derive(Debug, Clone)]
pub struct AccumulatorInstance<E: PairingEngine> {
    /// The polynomial commitment.
    pub commitment: E::G1Affine,

    /// The claimed evaluation.
    pub evaluation: E::Fr,

    /// The evaluation point.
    pub point: E::Fr,

    /// The quotient polynomial commitment (opening proof).
    pub quotient: E::G1Affine,
}

/// Accumulator for folding multiple proofs.
pub struct ProofAccumulator<E: PairingEngine> {
    /// Current accumulated state.
    instances: Vec<AccumulatorInstance<E>>,

    /// Transcript for generating challenges.
    transcript: Transcript,

    _engine: PhantomData<E>,
}

impl<E: PairingEngine> ProofAccumulator<E> {
    /// Create a new empty accumulator.
    pub fn new(domain: &str) -> Self {
        Self {
            instances: Vec::new(),
            transcript: Transcript::new(domain),
            _engine: PhantomData,
        }
    }

    /// Add an instance to the accumulator.
    pub fn add(&mut self, instance: AccumulatorInstance<E>) {
        // Add instance to transcript for Fiat-Shamir
        self.transcript.append_g1::<E>("commitment", &instance.commitment);
        self.transcript.append_scalar::<E>("evaluation", &instance.evaluation);
        self.transcript.append_scalar::<E>("point", &instance.point);

        self.instances.push(instance);
    }

    /// Get the number of accumulated instances.
    pub fn len(&self) -> usize {
        self.instances.len()
    }

    /// Check if the accumulator is empty.
    pub fn is_empty(&self) -> bool {
        self.instances.is_empty()
    }

    /// Fold all instances into a single accumulated proof.
    ///
    /// Uses random linear combination with Fiat-Shamir challenges.
    /// Implements Batched KZG:
    /// C_acc = Σ r_i * (C_i - [v_i]_1 + z_i * π_i)
    /// π_acc = Σ r_i * π_i
    pub fn fold(&mut self) -> Result<AccumulatedProof<E>, String> {
        use group::{Curve, Group};

        if self.instances.is_empty() {
            return Err("No instances to fold".to_string());
        }

        // Generate random challenges for linear combination
        // Note: For security, challenges must be sampled after seeing all instances
        let mut challenges: Vec<E::Fr> = Vec::with_capacity(self.instances.len());
        for i in 0..self.instances.len() {
            // We include the index in the label to ensure unique challenges
            let challenge = self.transcript.challenge_scalar::<E>(&format!("fold_{}", i));
            challenges.push(challenge);
        }

        let mut adjusted_commitment_acc = E::G1::identity();
        let mut combined_quotient_acc = E::G1::identity();

        let g1_gen = E::G1::generator();

        for (instance, r) in self.instances.iter().zip(challenges.iter()) {
            let comm_g1: E::G1 = instance.commitment.into();
            let quot_g1: E::G1 = instance.quotient.into();

            // [v]₁ = v * G1
            let v_g1 = g1_gen * instance.evaluation;

            // Adjusted commitment part: r * (C - [v] + z*π)
            let adjusted = (comm_g1 - v_g1 + quot_g1 * instance.point) * *r;
            adjusted_commitment_acc += adjusted;

            // Combined quotient part: r * π
            combined_quotient_acc += quot_g1 * *r;
        }

        let count = self.instances.len();
        self.instances.clear();

        Ok(AccumulatedProof {
            adjusted_commitment: adjusted_commitment_acc.to_affine(),
            combined_quotient: combined_quotient_acc.to_affine(),
            count,
        })
    }

    /// Verify that an accumulated proof is valid.
    ///
    /// Checks: e(adjusted_commitment, [1]₂) == e(combined_quotient, [τ]₂)
    ///
    /// # Optimization (Dan Boneh recommendation)
    ///
    /// We use a single `multi_pairing` call instead of two separate pairings.
    /// This leverages the shared final exponentiation for ~40% faster verification.
    ///
    /// The pairing equation is rearranged:
    /// - Original: e(A, G₂) == e(Q, τG₂)
    /// - Rearranged: e(A, G₂) · e(-Q, τG₂) == 1 (identity in Gt)
    ///
    /// By negating Q in G1, we can check if the product equals the identity element.
    pub fn verify_accumulated(
        proof: &AccumulatedProof<E>,
        srs_g2: &E::G2Affine,
        tau_g2: &E::G2Affine,
    ) -> bool {
        use group::{Curve, Group};

        // Negate the combined quotient: -Q in G1
        let combined_q: E::G1 = proof.combined_quotient.into();
        let neg_quotient = (-combined_q).to_affine();

        // Multi-pairing check: e(A, G₂) · e(-Q, τG₂) == 1
        // This is more efficient than two separate pairings due to shared final exp.
        let result = E::multi_pairing([
            (&proof.adjusted_commitment, srs_g2),
            (&neg_quotient, tau_g2),
        ]);

        result == E::Gt::identity()
    }

    /// Verify accumulated proof using separate pairings (for comparison/testing).
    /// 
    /// This is the original implementation kept for testing equivalence.
    #[cfg(test)]
    pub fn verify_accumulated_separate_pairings(
        proof: &AccumulatedProof<E>,
        srs_g2: &E::G2Affine,
        tau_g2: &E::G2Affine,
    ) -> bool {
        let lhs = E::pairing(&proof.adjusted_commitment, srs_g2);
        let rhs = E::pairing(&proof.combined_quotient, tau_g2);
        lhs == rhs
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::backend::bn254::Bn254;
    use ff::Field;
    use group::{Curve, Group};
    use halo2curves::bn256::{Fr, G1};
    use rand::rngs::OsRng;

    fn mock_instance() -> AccumulatorInstance<Bn254> {
        AccumulatorInstance {
            commitment: G1::random(OsRng).to_affine(),
            evaluation: Fr::random(OsRng),
            point: Fr::random(OsRng),
            quotient: G1::random(OsRng).to_affine(),
        }
    }

    // TDD tests

    #[test]
    fn accumulator_starts_empty() {
        let acc = ProofAccumulator::<Bn254>::new("test");
        assert!(acc.is_empty());
        assert_eq!(acc.len(), 0);
    }

    #[test]
    fn accumulator_tracks_instances() {
        let mut acc = ProofAccumulator::<Bn254>::new("test");

        acc.add(mock_instance());
        acc.add(mock_instance());

        assert_eq!(acc.len(), 2);
        assert!(!acc.is_empty());
    }

    #[test]
    fn fold_empty_fails() {
        let mut acc = ProofAccumulator::<Bn254>::new("test");

        let result = acc.fold();
        assert!(result.is_err());
    }

    #[test]
    fn fold_multiple_produces_valid_accumulation() {
        let mut acc = ProofAccumulator::<Bn254>::new("test");

        acc.add(mock_instance());
        acc.add(mock_instance());
        acc.add(mock_instance());

        let folded = acc.fold().unwrap();

        assert_eq!(folded.count, 3);
        // Adjusted commitment should not be identity (with overwhelming probability)
        assert_ne!(folded.adjusted_commitment, G1::identity().to_affine());
        // Combined quotient should not be identity
        assert_ne!(folded.combined_quotient, G1::identity().to_affine());
    }

    #[test]
    fn fold_clears_instances() {
        let mut acc = ProofAccumulator::<Bn254>::new("test");

        acc.add(mock_instance());
        acc.add(mock_instance());

        let _ = acc.fold().unwrap();

        assert!(acc.is_empty());
    }

    #[test]
    fn fold_is_deterministic() {
        // Two accumulators with same instances should produce same result
        let instance1 = AccumulatorInstance {
            commitment: G1::generator().to_affine(),
            evaluation: Fr::from(42u64),
            point: Fr::from(7u64),
            quotient: (G1::generator() * Fr::from(3u64)).to_affine(),
        };
        let instance2 = AccumulatorInstance {
            commitment: (G1::generator() * Fr::from(2u64)).to_affine(),
            evaluation: Fr::from(100u64),
            point: Fr::from(13u64),
            quotient: (G1::generator() * Fr::from(5u64)).to_affine(),
        };

        let mut acc1 = ProofAccumulator::<Bn254>::new("test");
        acc1.add(instance1.clone());
        acc1.add(instance2.clone());

        let mut acc2 = ProofAccumulator::<Bn254>::new("test");
        acc2.add(instance1);
        acc2.add(instance2);

        let folded1 = acc1.fold().unwrap();
        let folded2 = acc2.fold().unwrap();

        assert_eq!(folded1.adjusted_commitment, folded2.adjusted_commitment);
        assert_eq!(folded1.combined_quotient, folded2.combined_quotient);
    }

    #[test]
    fn multi_pairing_and_separate_pairing_are_equivalent() {
        // Create deterministic instances
        let instance1 = AccumulatorInstance {
            commitment: G1::generator().to_affine(),
            evaluation: Fr::from(42u64),
            point: Fr::from(7u64),
            quotient: (G1::generator() * Fr::from(3u64)).to_affine(),
        };
        let instance2 = AccumulatorInstance {
            commitment: (G1::generator() * Fr::from(2u64)).to_affine(),
            evaluation: Fr::from(100u64),
            point: Fr::from(13u64),
            quotient: (G1::generator() * Fr::from(5u64)).to_affine(),
        };

        let mut acc = ProofAccumulator::<Bn254>::new("equivalence_test");
        acc.add(instance1);
        acc.add(instance2);

        let folded = acc.fold().unwrap();

        // Create SRS-like G2 points for testing
        let g2_gen = halo2curves::bn256::G2::generator().to_affine();
        let tau = Fr::from(0xDEADBEEFu64);
        let tau_g2 = (halo2curves::bn256::G2::generator() * tau).to_affine();

        // Both implementations should return the same result
        let result_multi = ProofAccumulator::<Bn254>::verify_accumulated(&folded, &g2_gen, &tau_g2);
        let result_separate = ProofAccumulator::<Bn254>::verify_accumulated_separate_pairings(&folded, &g2_gen, &tau_g2);

        assert_eq!(result_multi, result_separate, 
            "multi_pairing and separate pairing should give same result");
    }
}
