samaharam 0.1.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 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,
        })
    }

    /// 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))
    }

    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);
    }
}