samaharam 0.2.0

//! Groth16 compatibility layer.
//!
//! Provides conversion between samaharam's native proof format and
//! the standard Groth16 proof format used by most EVM tooling
//! (snarkjs, circom, etc.).

use crate::traits::PairingEngine;

/// Standard serialized size for Groth16 proofs.
/// G1 (64 bytes uncompressed) + G2 (128 bytes uncompressed) + G1 (64 bytes)
pub const GROTH16_PROOF_SIZE: usize = 256;

/// Standard Groth16 proof format with [A, B, C] structure.
///
/// This format is compatible with:
/// - snarkjs
/// - circom
/// - Iden3's rapidsnark
/// - Most EVM verifier contracts
#[derive(Debug, Clone)]
pub struct Groth16Proof<E: PairingEngine> {
    /// A ∈ G1 - first pairing element
    pub a: E::G1Affine,

    /// B ∈ G2 - second pairing element (stored as raw bytes for serialization)
    pub b_bytes: [u8; 128],

    /// C ∈ G1 - third pairing element
    pub c: E::G1Affine,
}

impl<E: PairingEngine> Groth16Proof<E> {
    /// Create from G1 points with G2 bytes.
    pub fn new(a: E::G1Affine, b_bytes: [u8; 128], c: E::G1Affine) -> Self {
        Self { a, b_bytes, c }
    }

    /// Create from all components using GroupEncoding for serialization.
    pub fn from_components(a: E::G1Affine, b: &E::G2Affine, c: E::G1Affine) -> Self {
        use group::GroupEncoding;

        let b_encoded = b.to_bytes();
        let b_ref: &[u8] = b_encoded.as_ref();
        let mut b_bytes = [0u8; 128];
        let len = b_ref.len().min(128);
        b_bytes[..len].copy_from_slice(&b_ref[..len]);

        Self { a, b_bytes, c }
    }

    /// Serialize to standard ABI-encoded bytes.
    ///
    /// Format: `[A.x, A.y, B.x[0], B.x[1], B.y[0], B.y[1], C.x, C.y]`
    /// Each coordinate is 32 bytes (big-endian uint256).
    pub fn to_abi_bytes(&self) -> Vec<u8> {
        use group::GroupEncoding;

        let mut bytes = Vec::with_capacity(GROTH16_PROOF_SIZE);

        // Serialize A (G1): 64 bytes
        let a_bytes = self.a.to_bytes();
        let a_ref: &[u8] = a_bytes.as_ref();
        bytes.extend_from_slice(a_ref);
        // Pad to 64 bytes
        bytes.resize(64, 0);

        // Serialize B (G2): 128 bytes (already stored as bytes)
        bytes.extend_from_slice(&self.b_bytes);

        // Serialize C (G1): 64 bytes
        let c_bytes = self.c.to_bytes();
        let c_ref: &[u8] = c_bytes.as_ref();
        bytes.extend_from_slice(c_ref);
        // Pad to 256 total
        bytes.resize(GROTH16_PROOF_SIZE, 0);

        bytes
    }

    /// Deserialize A and C points from ABI-encoded bytes.
    /// B is stored as raw bytes since G2 serialization varies.
    pub fn from_abi_bytes(data: &[u8]) -> Result<Self, String> {
        if data.len() < GROTH16_PROOF_SIZE {
            return Err(format!(
                "proof data too short: {} < {}",
                data.len(),
                GROTH16_PROOF_SIZE
            ));
        }

        use group::GroupEncoding;

        // Parse A (G1)
        let mut a_repr = <E::G1Affine as GroupEncoding>::Repr::default();
        let a_slice: &mut [u8] = a_repr.as_mut();
        let a_len = a_slice.len().min(64);
        a_slice[..a_len].copy_from_slice(&data[..a_len]);
        let a = E::G1Affine::from_bytes(&a_repr);
        if a.is_none().into() {
            return Err("invalid A point".to_string());
        }

        // Store B bytes directly
        let mut b_bytes = [0u8; 128];
        b_bytes.copy_from_slice(&data[64..192]);

        // Parse C (G1)
        let mut c_repr = <E::G1Affine as GroupEncoding>::Repr::default();
        let c_slice: &mut [u8] = c_repr.as_mut();
        let c_len = c_slice.len().min(64);
        c_slice[..c_len].copy_from_slice(&data[192..192 + c_len]);
        let c = E::G1Affine::from_bytes(&c_repr);
        if c.is_none().into() {
            return Err("invalid C point".to_string());
        }

        Ok(Self {
            a: a.unwrap(),
            b_bytes,
            c: c.unwrap(),
        })
    }

    /// Convert to snarkjs-compatible JSON proof format.
    pub fn to_snarkjs_json(&self) -> String {
        use group::GroupEncoding;

        let a_bytes = self.a.to_bytes();
        let c_bytes = self.c.to_bytes();

        // Format as snarkjs expects (array of strings)
        format!(
            r#"{{
  "pi_a": ["0x{}", "0x{}", "1"],
  "pi_b": [["0x{}", "0x{}"], ["0x{}", "0x{}"], ["1", "0"]],
  "pi_c": ["0x{}", "0x{}", "1"],
  "protocol": "groth16"
}}"#,
            hex::encode(&a_bytes.as_ref()[..32.min(a_bytes.as_ref().len())]),
            hex::encode(
                &a_bytes.as_ref()[32.min(a_bytes.as_ref().len())..64.min(a_bytes.as_ref().len())]
            ),
            hex::encode(&self.b_bytes[..32]),
            hex::encode(&self.b_bytes[32..64]),
            hex::encode(&self.b_bytes[64..96]),
            hex::encode(&self.b_bytes[96..128]),
            hex::encode(&c_bytes.as_ref()[..32.min(c_bytes.as_ref().len())]),
            hex::encode(
                &c_bytes.as_ref()[32.min(c_bytes.as_ref().len())..64.min(c_bytes.as_ref().len())]
            ),
        )
    }

    /// Convert to Solidity calldata format.
    ///
    /// Returns a tuple (a, b, c) ready for Solidity verifier:
    /// ```solidity
    /// function verify(
    ///     uint256[2] memory a,
    ///     uint256[2][2] memory b,
    ///     uint256[2] memory c,
    ///     uint256[] memory publicInputs
    /// ) external returns (bool)
    /// ```
    pub fn to_solidity_calldata(&self) -> SolidityCalldata {
        use group::GroupEncoding;

        let a = self.a.to_bytes();
        let c = self.c.to_bytes();

        // Extract coordinates
        let a_ref: &[u8] = a.as_ref();
        let c_ref: &[u8] = c.as_ref();

        SolidityCalldata {
            a: [
                bytes_to_u256(&a_ref[..32.min(a_ref.len())]),
                bytes_to_u256(&a_ref[32.min(a_ref.len())..64.min(a_ref.len())]),
            ],
            b: [
                [
                    bytes_to_u256(&self.b_bytes[..32]),
                    bytes_to_u256(&self.b_bytes[32..64]),
                ],
                [
                    bytes_to_u256(&self.b_bytes[64..96]),
                    bytes_to_u256(&self.b_bytes[96..128]),
                ],
            ],
            c: [
                bytes_to_u256(&c_ref[..32.min(c_ref.len())]),
                bytes_to_u256(&c_ref[32.min(c_ref.len())..64.min(c_ref.len())]),
            ],
        }
    }
}

/// Solidity-compatible calldata representation.
#[derive(Debug, Clone)]
pub struct SolidityCalldata {
    /// A point coordinates [x, y]
    pub a: [String; 2],
    /// B point coordinates [[x0, x1], [y0, y1]]
    pub b: [[String; 2]; 2],
    /// C point coordinates [x, y]
    pub c: [String; 2],
}

impl SolidityCalldata {
    /// Format as Solidity function argument.
    pub fn to_solidity_args(&self) -> String {
        format!(
            "[{}, {}], [[{}, {}], [{}, {}]], [{}, {}]",
            self.a[0],
            self.a[1],
            self.b[0][0],
            self.b[0][1],
            self.b[1][0],
            self.b[1][1],
            self.c[0],
            self.c[1]
        )
    }
}

/// Convert bytes to uint256 string representation.
fn bytes_to_u256(bytes: &[u8]) -> String {
    if bytes.is_empty() {
        return "0".to_string();
    }

    // Pad to 32 bytes
    let mut padded = [0u8; 32];
    let len = bytes.len().min(32);
    padded[32 - len..].copy_from_slice(&bytes[..len]);

    format!("0x{}", hex::encode(padded))
}

/// Convert aggregated proof to Groth16 format.
///
/// This converts samaharam's accumulated proof structure to
/// the standard Groth16 [A, B, C] format for EVM verification.
pub fn convert_to_groth16<E: PairingEngine>(
    aggregated_data: &[u8],
    srs_g2: &E::G2Affine,
) -> Result<Groth16Proof<E>, String> {
    use group::GroupEncoding;

    // Minimum size: adjusted_commitment (48) + combined_quotient (48) + count (4)
    if aggregated_data.len() < 100 {
        return Err("aggregated data too short".to_string());
    }

    // Parse adjusted commitment (G1) -> becomes A
    let mut a_repr = <E::G1Affine as GroupEncoding>::Repr::default();
    let a_slice: &mut [u8] = a_repr.as_mut();
    let len = a_slice.len().min(48);
    a_slice[..len].copy_from_slice(&aggregated_data[..len]);
    let a = E::G1Affine::from_bytes(&a_repr);

    // Parse combined quotient (G1) -> becomes C
    let mut c_repr = <E::G1Affine as GroupEncoding>::Repr::default();
    let c_slice: &mut [u8] = c_repr.as_mut();
    c_slice[..len].copy_from_slice(&aggregated_data[48..48 + len]);
    let c = E::G1Affine::from_bytes(&c_repr);

    if a.is_none().into() || c.is_none().into() {
        return Err("invalid aggregated proof points".to_string());
    }

    // Serialize G2 point for B
    let b_bytes_encoded = srs_g2.to_bytes();
    let b_ref: &[u8] = b_bytes_encoded.as_ref();
    let mut b_bytes = [0u8; 128];
    let b_len = b_ref.len().min(128);
    b_bytes[..b_len].copy_from_slice(&b_ref[..b_len]);

    Ok(Groth16Proof {
        a: a.unwrap(),
        b_bytes,
        c: c.unwrap(),
    })
}

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

    #[test]
    fn groth16_proof_serialization_roundtrip() {
        use group::GroupEncoding;

        let g2 = G2::random(OsRng).to_affine();
        let g2_bytes = g2.to_bytes();
        let g2_ref: &[u8] = g2_bytes.as_ref();

        let mut b_bytes = [0u8; 128];
        let len = g2_ref.len().min(128);
        b_bytes[..len].copy_from_slice(&g2_ref[..len]);

        let proof = Groth16Proof::<Bn254> {
            a: G1::random(OsRng).to_affine(),
            b_bytes,
            c: G1::random(OsRng).to_affine(),
        };

        let bytes = proof.to_abi_bytes();
        assert_eq!(bytes.len(), GROTH16_PROOF_SIZE);

        let decoded = Groth16Proof::<Bn254>::from_abi_bytes(&bytes);
        assert!(decoded.is_ok());
    }

    #[test]
    fn groth16_proof_to_snarkjs_json() {
        let proof = Groth16Proof::<Bn254> {
            a: G1::generator().to_affine(),
            b_bytes: [0u8; 128],
            c: G1::generator().to_affine(),
        };

        let json = proof.to_snarkjs_json();

        assert!(json.contains("pi_a"));
        assert!(json.contains("pi_b"));
        assert!(json.contains("pi_c"));
        assert!(json.contains("groth16"));
    }

    #[test]
    fn solidity_calldata_format() {
        let proof = Groth16Proof::<Bn254> {
            a: G1::generator().to_affine(),
            b_bytes: [0u8; 128],
            c: G1::generator().to_affine(),
        };

        let calldata = proof.to_solidity_calldata();

        assert!(calldata.a[0].starts_with("0x"));
        assert!(calldata.b[0][0].starts_with("0x"));

        let args = calldata.to_solidity_args();
        assert!(args.contains("0x"));
    }

    #[test]
    fn bytes_to_u256_pads_correctly() {
        let bytes = vec![0x01, 0x02];
        let result = bytes_to_u256(&bytes);

        // Should be padded to 32 bytes with leading zeros
        assert!(result.starts_with("0x"));
        assert_eq!(result.len(), 66); // 0x + 64 hex chars
    }

    #[test]
    fn from_components_creates_valid_proof() {
        use group::GroupEncoding;

        let a = G1::random(OsRng).to_affine();
        let b = G2::random(OsRng).to_affine();
        let c = G1::random(OsRng).to_affine();

        let proof = Groth16Proof::<Bn254>::from_components(a.clone(), &b, c.clone());

        assert_eq!(proof.a, a);
        assert_eq!(proof.c, c);
        // b_bytes should contain the serialized G2 point
        assert!(!proof.b_bytes.iter().all(|&x| x == 0));
    }
}