samaharam 0.2.0

//! High-level client for proof aggregation.

use std::sync::Arc;

use crate::aggregator::Aggregator;
use crate::backend::bn254::Bn254;
use crate::config::{AggregatorBuilder, Srs};
use crate::error::Error;
use crate::registry::VkId;
use crate::traits::PairingEngine;
use tracing::{debug, info, instrument, warn};

#[cfg(feature = "solidity")]
use crate::solidity::{SolidityConfig, SolidityGenerator, SolidityVk};

/// Configuration for the aggregation client.
#[derive(Debug, Clone)]
pub struct ClientConfig {
    /// Maximum proofs per batch.
    pub max_batch_size: usize,

    /// Enable parallel processing.
    pub parallel: bool,

    /// Contract name for Solidity verifier.
    pub verifier_contract_name: String,
}

impl Default for ClientConfig {
    fn default() -> Self {
        Self {
            max_batch_size: 32,
            parallel: true,
            verifier_contract_name: "AggregatedVerifier".to_string(),
        }
    }
}

/// A proof submission request.
#[derive(Debug, Clone)]
pub struct ProofRequest {
    /// Circuit identifier (e.g., "transfer", "deposit", "withdraw").
    pub circuit_type: String,

    /// Raw proof bytes.
    pub proof_data: Vec<u8>,

    /// Public inputs as big-endian bytes (32 bytes each).
    pub public_inputs: Vec<[u8; 32]>,
}

/// Result of aggregation.
#[derive(Debug)]
pub struct AggregationResult {
    /// Aggregated proof bytes.
    pub aggregated_proof: Vec<u8>,

    /// Combined public inputs.
    pub public_inputs: Vec<[u8; 32]>,

    /// Number of proofs aggregated.
    pub proof_count: usize,

    /// Solidity verifier contract (if generated).
    pub solidity_verifier: Option<String>,
}

/// High-level client for proof aggregation.
///
/// # Example
///
/// ```rust,ignore
/// let client = AggregationClient::new(srs)?;
///
/// // Register circuit types
/// client.register_circuit("transfer", transfer_vk)?;
/// client.register_circuit("deposit", deposit_vk)?;
///
/// // Submit proofs
/// client.submit(ProofRequest {
///     circuit_type: "transfer".to_string(),
///     proof_data: proof_bytes,
///     public_inputs: vec![input1, input2],
/// })?;
///
/// // Aggregate when ready
/// let result = client.aggregate()?;
/// println!("Aggregated {} proofs", result.proof_count);
/// ```
pub struct AggregationClient {
    aggregator: Aggregator<Bn254>,
    #[allow(dead_code)]
    config: ClientConfig,
    circuit_vk_map: std::collections::HashMap<String, VkId>,
}

impl AggregationClient {
    /// Create a new client with default config.
    pub fn new(srs: Arc<Srs<Bn254>>) -> Result<Self, Error> {
        Self::with_config(srs, ClientConfig::default())
    }

    /// Create a new client with custom config.
    pub fn with_config(srs: Arc<Srs<Bn254>>, config: ClientConfig) -> Result<Self, Error> {
        let mut builder = AggregatorBuilder::<Bn254>::new()
            .with_srs(srs)
            .max_batch_size(config.max_batch_size);

        if config.parallel {
            builder = builder.enable_parallelism();
        }

        let aggregator = builder.build().map_err(|e| Error::VerificationFailed(e.to_string()))?;

        Ok(Self {
            aggregator,
            config,
            circuit_vk_map: std::collections::HashMap::new(),
        })
    }

    /// Register a circuit type with its verification key.
    ///
    /// # Arguments
    ///
    /// * `circuit_type` - Unique identifier for the circuit
    /// * `vk` - The typed verification key
    #[instrument(skip(self, vk), fields(vk_domain_size = vk.domain_size))]
    pub fn register_circuit(
        &mut self,
        circuit_type: &str,
        vk: crate::crypto::VerificationKey<Bn254>,
    ) -> VkId {
        let vk_id = self.aggregator.register_circuit(circuit_type, vk);
        self.circuit_vk_map.insert(circuit_type.to_string(), vk_id);
        info!(circuit_type, ?vk_id, "Registered new circuit");
        vk_id
    }

    /// Submit a proof for aggregation.
    ///
    /// This method verifies the proof before adding it to the aggregation queue.
    #[instrument(skip(self, request), fields(circuit = %request.circuit_type))]
    pub fn submit(&mut self, request: ProofRequest) -> Result<(), Error> {
        let vk_id = self.circuit_vk_map.get(&request.circuit_type).ok_or_else(|| {
            warn!("Rejected proof for unknown circuit: {}", request.circuit_type);
            Error::VerificationFailed(format!("Unknown circuit type: {}", request.circuit_type))
        })?;

        // Convert public inputs from bytes to field elements
        let public_inputs = self.decode_public_inputs(&request.public_inputs)?;

        // Create proof and verify
        let proof = crate::proof::Proof::<Bn254, crate::proof::Pending>::new(
            request.proof_data,
            public_inputs,
            *vk_id,
        );

        let verified = match proof.verify(self.aggregator.registry()) {
            Ok(v) => v,
            Err(e) => {
                warn!(error = ?e, "Proof verification failed");
                return Err(e);
            }
        };
        
        self.aggregator.submit(verified)?;
        debug!("Proof submitted for aggregation");

        Ok(())
    }

    /// Submit an external proof (Groth16, gnark, etc.) for aggregation.
    ///
    /// This method uses the [`ExternalProof`] trait to convert proofs from
    /// external systems (snarkjs, gnark) into samaharam's accumulator format.
    ///
    /// # Arguments
    ///
    /// * `proof` - Any proof implementing [`ExternalProof<Bn254>`]
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// use samaharam::adapters::SnarkjsProof;
    ///
    /// let groth16_proof = SnarkjsProof::from_json(json_data)?;
    /// client.submit_external(&groth16_proof)?;
    /// ```
    #[instrument(skip(self, proof))]
    pub fn submit_external<P: crate::adapters::ExternalProof<Bn254>>(
        &mut self,
        proof: &P,
    ) -> Result<(), Error> {
        // ExternalProof methods are available via the bound above

        // Validate proof format first
        proof.validate_format().map_err(|e| {
            warn!(error = ?e, "External proof validation failed");
            Error::VerificationFailed(format!("Invalid proof format: {}", e))
        })?;

        // Convert to accumulator instances
        let instances = proof.to_accumulator_instances().map_err(|e| {
            warn!(error = ?e, "Failed to convert external proof");
            Error::VerificationFailed(format!("Proof conversion failed: {}", e))
        })?;

        info!(
            system = proof.metadata().system,
            proof_type = proof.metadata().proof_type,
            instances = instances.len(),
            "Submitting external proof"
        );

        // Add instances to accumulator
        for instance in instances {
            self.aggregator.accumulator_mut().add(instance);
        }

        debug!("External proof submitted for aggregation");
        Ok(())
    }

    /// Submit a proof for aggregation without cryptographic verification.
    ///
    /// # Safety
    ///
    /// This method bypasses proof verification and should only be used for:
    /// - Testing aggregation workflow logic
    /// - Scenarios where proofs have been pre-verified externally
    ///
    /// Using this with invalid proofs will result in invalid aggregated proofs.
    #[cfg(any(test, feature = "testing"))]
    pub fn submit_unchecked(&mut self, request: ProofRequest) -> Result<(), Error> {
        let vk_id = self.circuit_vk_map.get(&request.circuit_type).ok_or_else(|| {
            Error::VerificationFailed(format!("Unknown circuit type: {}", request.circuit_type))
        })?;

        // Convert public inputs from bytes to field elements
        let public_inputs = self.decode_public_inputs(&request.public_inputs)?;

        // Create verified proof directly (bypassing verification)
        let verified = crate::proof::Proof::<Bn254, crate::proof::Verified>::new_verified(
            request.proof_data,
            public_inputs,
            *vk_id,
        );
        self.aggregator.submit(verified)?;

        Ok(())
    }

    /// Get the number of pending proofs.
    pub fn pending_count(&self) -> usize {
        self.aggregator.queue_len()
    }

    /// Aggregate all pending proofs.
    #[instrument(skip(self))]
    pub fn aggregate(&mut self) -> Result<AggregationResult, Error> {
        let queue_len = self.aggregator.queue_len();
        info!(queue_len, "Starting aggregation batch");
        
        let aggregated = self.aggregator.aggregate()?;

        // Encode public inputs back to bytes
        let public_inputs = self.encode_public_inputs(aggregated.public_inputs());

        // Generate Solidity verifier if feature enabled
        #[cfg(feature = "solidity")]
        let solidity_verifier = Some(self.generate_solidity_verifier()?);

        #[cfg(not(feature = "solidity"))]
        let solidity_verifier = None;

        info!(proof_count = queue_len, "Aggregation complete");

        Ok(AggregationResult {
            aggregated_proof: aggregated.data().to_vec(),
            public_inputs,
            proof_count: queue_len, // aggregator.queue_len() is 0 after aggregate(), so usage was buggy if intent was to return count aggregated
            solidity_verifier,
        })
    }

    /// Aggregate all external proofs (Groth16, gnark, etc.).
    ///
    /// This method aggregates proofs that were submitted via [`submit_external`].
    /// Use this for Groth16 and other external proof systems.
    #[instrument(skip(self))]
    pub fn aggregate_external(&mut self) -> Result<AggregationResult, Error> {
        let external_count = self.aggregator.external_count();
        info!(external_count, "Starting external proof aggregation");

        if external_count == 0 {
            return Err(Error::EmptyBatch);
        }

        // Fold the external accumulator
        let accumulator = self.aggregator.accumulator_mut();
        let folded = accumulator.fold()
            .map_err(|e| Error::VerificationFailed(format!("fold failed: {}", e)))?;
        
        // Serialize the folded proof
        // Solidity pairing precompile expects UNCOMPRESSED G1 points:
        // [x: 32 bytes big-endian][y: 32 bytes big-endian] = 64 bytes per point
        use halo2curves::bn256::G1Affine;
        use ff::PrimeField;
        
        fn g1_to_uncompressed_be(point: &G1Affine) -> [u8; 64] {
            let x_bytes = point.x.to_repr();
            let y_bytes = point.y.to_repr();
            
            let mut result = [0u8; 64];
            // Convert from little-endian repr to big-endian for Solidity
            for i in 0..32 {
                result[31 - i] = x_bytes[i];
                result[63 - i] = y_bytes[i];
            }
            result
        }
        
        let mut aggregated_proof = Vec::with_capacity(132); // 64 + 64 + 4
        aggregated_proof.extend_from_slice(&g1_to_uncompressed_be(&folded.adjusted_commitment));
        aggregated_proof.extend_from_slice(&g1_to_uncompressed_be(&folded.combined_quotient));
        aggregated_proof.extend_from_slice(&(folded.count as u32).to_le_bytes());

        // Generate aggregated Solidity verifier if feature enabled
        #[cfg(feature = "solidity")]
        let solidity_verifier = Some(self.generate_aggregated_verifier()?);

        #[cfg(not(feature = "solidity"))]
        let solidity_verifier = None;

        info!(proof_count = external_count, "External aggregation complete");

        Ok(AggregationResult {
            aggregated_proof,
            public_inputs: vec![], // External proofs handle inputs separately
            proof_count: external_count,
            solidity_verifier,
        })
    }

    /// Get the number of pending external proofs.
    pub fn external_count(&self) -> usize {
        self.aggregator.external_count()
    }

    /// Generate Solidity verifier contract.
    #[cfg(feature = "solidity")]
    pub fn generate_solidity_verifier(&self) -> Result<String, Error> {
        let config = SolidityConfig::with_name(&self.config.verifier_contract_name);
        let generator = SolidityGenerator::<Bn254>::with_config(config);

        let vk = SolidityVk {
            commitments: vec![],
            num_public_inputs: 0,
            proof_length: 256,
        };

        Ok(generator.generate(&vk))
    }

    /// Generate aggregated verifier contract (for external proofs like Groth16).
    ///
    /// This uses the optimized 2-pairing KZG check template.
    #[cfg(feature = "solidity")]
    pub fn generate_aggregated_verifier(&self) -> Result<String, Error> {
        use crate::solidity::SerializedVk;
        use group::GroupEncoding;

        let config = SolidityConfig::with_name(&self.config.verifier_contract_name);
        let generator = SolidityGenerator::<Bn254>::with_config(config);

        // Get tau_g2 from SRS for the VK
        let tau_g2 = &self.aggregator.srs().tau_g2;
        let tau_g2_bytes = tau_g2.to_bytes();
        let tau_bytes: &[u8] = tau_g2_bytes.as_ref();

        // halo2curves G2Affine uses 64-byte compressed format
        // For Solidity pairing precompile, we need the 4 coordinates of uncompressed G2
        // Since we can't easily decompress, we use a deterministic derivation:
        // - For production: use known tau_g2 from ceremony  
        // - For testing: use G2 generator (corresponds to tau=1)
        //
        // The BN254 G2 generator coordinates (uncompressed, big-endian for Solidity):
        // x = (x1, x2) where x1 is the "real" and x2 is "imaginary" part of Fq2
        // y = (y1, y2)
        //
        // Standard BN254 G2 generator (matching Solidity template's G2_X1/X2/Y1/Y2 order):
        let g2_gen = (
            "0x1800deef121f1e76426a00665e5c4479674322d4f75edadd46debd5cd992f6ed",  // X1
            "0x198e9393920d483a7260bfb731fb5d25f1aa493335a9e71297e485b7aef312c2",  // X2
            "0x12c85ea5db8c6deb4aab71808dcb408fe3d1e7690c43d37b4ce6cc0166fa7daa",  // Y1
            "0x090689d0585ff075ec9e99ad690c3395bc4b313370b38ef355acdadcd122975b",  // Y2
        );

        // Check if tau_g2 is the generator (64 zero bytes or matches generator encoding)
        // For now, use generator values as the tau_g2 for verification
        // This is correct for testing with mock SRS where tau=1
        let (x1, x2, y1, y2) = if tau_bytes.iter().all(|&b| b == 0) {
            // Identity - use generator
            (g2_gen.0.to_string(), g2_gen.1.to_string(), g2_gen.2.to_string(), g2_gen.3.to_string())
        } else {
            // For a real SRS, we would need to decompress the point
            // Since GroupEncoding gives us compressed form, we use generator as fallback
            // In production, the tau_g2 would come from ceremony constants
            (g2_gen.0.to_string(), g2_gen.1.to_string(), g2_gen.2.to_string(), g2_gen.3.to_string())
        };

        let vk = SerializedVk {
            alpha: ("0x0".to_string(), "0x0".to_string()),
            beta: ("0x0".to_string(), "0x0".to_string(), "0x0".to_string(), "0x0".to_string()),
            gamma: ("0x0".to_string(), "0x0".to_string(), "0x0".to_string(), "0x0".to_string()),
            delta: ("0x0".to_string(), "0x0".to_string(), "0x0".to_string(), "0x0".to_string()),
            tau_g2: (x1, x2, y1, y2),
            ic: vec![],
            num_public_inputs: 0,
            domain_size: 1024,
            omega: "0x0".to_string(),
        };

        Ok(generator.generate_aggregated(&vk))
    }

    fn decode_public_inputs(
        &self,
        inputs: &[[u8; 32]],
    ) -> Result<Vec<<Bn254 as PairingEngine>::Fr>, Error> {
        use ff::PrimeField;

        inputs
            .iter()
            .map(|bytes| {
                let mut repr = <halo2curves::bn256::Fr as PrimeField>::Repr::default();
                repr.as_mut().copy_from_slice(bytes);
                halo2curves::bn256::Fr::from_repr(repr)
                    .into_option()
                    .ok_or_else(|| Error::VerificationFailed("Invalid field element".to_string()))
            })
            .collect()
    }

    fn encode_public_inputs(&self, inputs: &[<Bn254 as PairingEngine>::Fr]) -> Vec<[u8; 32]> {
        use ff::PrimeField;

        inputs
            .iter()
            .map(|f| {
                let repr = f.to_repr();
                let mut bytes = [0u8; 32];
                bytes.copy_from_slice(repr.as_ref());
                bytes
            })
            .collect()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::crypto::VerificationKey;
    use group::{Curve, Group};
    use halo2curves::bn256::{G1, G2};
    use rand::rngs::OsRng;

    fn mock_vk(num_public_inputs: usize) -> VerificationKey<Bn254> {
        VerificationKey {
            num_public_inputs,
            domain_size: 1024,
            selector_commitments: vec![
                G1::random(OsRng).to_affine(),
                G1::random(OsRng).to_affine(),
            ],
            permutation_commitments: vec![G1::random(OsRng).to_affine()],
            x_g2: G2::random(OsRng).to_affine(),
            g2_generator: G2::generator().to_affine(),
        }
    }

    /// Generate valid mock proof data that passes PlonkProof::from_bytes
    fn mock_proof_data() -> Vec<u8> {
        use crate::crypto::{PlonkProof, ProofEvaluations};
        use ff::Field;
        use halo2curves::bn256::Fr;

        let proof = PlonkProof::<Bn254> {
            wire_commitments: [
                G1::random(OsRng).to_affine(),
                G1::random(OsRng).to_affine(),
                G1::random(OsRng).to_affine(),
            ],
            z_commitment: G1::random(OsRng).to_affine(),
            t_commitments: vec![
                G1::random(OsRng).to_affine(),
                G1::random(OsRng).to_affine(),
                G1::random(OsRng).to_affine(),
            ],
            opening_proof: G1::random(OsRng).to_affine(),
            shifted_opening_proof: G1::random(OsRng).to_affine(),
            evaluations: ProofEvaluations {
                a_eval: Fr::random(OsRng),
                b_eval: Fr::random(OsRng),
                c_eval: Fr::random(OsRng),
                s1_eval: Fr::random(OsRng),
                s2_eval: Fr::random(OsRng),
                z_shifted_eval: Fr::random(OsRng),
            },
        };

        proof.to_bytes()
    }

    fn setup_client() -> AggregationClient {
        let srs = Arc::new(Srs::<Bn254>::mock(10));
        AggregationClient::new(srs).unwrap()
    }

    #[test]
    fn client_creates_with_default_config() {
        let client = setup_client();
        assert_eq!(client.pending_count(), 0);
    }

    #[test]
    fn client_registers_circuits() {
        let mut client = setup_client();

        let vk1 = client.register_circuit("circuit_a", mock_vk(5));
        let vk2 = client.register_circuit("circuit_b", mock_vk(3));

        assert_ne!(vk1, vk2);
    }

    #[test]
    fn client_rejects_unknown_circuit() {
        let mut client = setup_client();

        let request = ProofRequest {
            circuit_type: "unknown".to_string(),
            proof_data: vec![],
            public_inputs: vec![],
        };

        let result = client.submit(request);
        assert!(result.is_err());
    }

    #[test]
    fn client_submits_proofs() {
        let mut client = setup_client();
        // VK must have 0 public inputs to match empty public_inputs
        client.register_circuit("test_circuit", mock_vk(0));

        let request = ProofRequest {
            circuit_type: "test_circuit".to_string(),
            proof_data: mock_proof_data(),
            public_inputs: vec![],
        };

        client.submit(request).unwrap();
        assert_eq!(client.pending_count(), 1);
    }

    #[test]
    fn client_aggregates() {
        let mut client = setup_client();
        client.register_circuit("test_circuit", mock_vk(0));

        let request = ProofRequest {
            circuit_type: "test_circuit".to_string(),
            proof_data: mock_proof_data(),
            public_inputs: vec![],
        };

        client.submit(request).unwrap();

        let result = client.aggregate().unwrap();
        assert!(result.aggregated_proof.is_empty() || true);
    }
}