samaharam 0.1.0

//! Point and field element parsing utilities for external proofs.
//!
//! BN254 curve point formats:
//! - G1 uncompressed: 64 bytes (32 x, 32 y)
//! - G1 compressed: 32 bytes (x with sign bit)
//! - G2 uncompressed: 128 bytes (32*4 for Fq2 x and y)

use halo2curves::bn256::{Fq, Fr, G1Affine, G2Affine};
use group::{Curve, GroupEncoding, prime::PrimeCurveAffine};

/// Parse a base field element (Fq) from 32 bytes (little-endian).
pub fn parse_fq_le(bytes: &[u8]) -> Result<Fq, String> {
    if bytes.len() != 32 {
        return Err(format!("Fq requires 32 bytes, got {}", bytes.len()));
    }

    let arr: [u8; 32] = bytes.try_into().map_err(|_| "Failed to convert bytes")?;
    let fq = Fq::from_bytes(&arr);
    
    if fq.is_none().into() {
        return Err("Invalid Fq field element bytes".to_string());
    }

    Ok(fq.unwrap())
}

/// Parse a base field element (Fq) from 32 bytes (big-endian).
pub fn parse_fq_be(bytes: &[u8]) -> Result<Fq, String> {
    if bytes.len() != 32 {
        return Err(format!("Fq requires 32 bytes, got {}", bytes.len()));
    }

    // Reverse to little-endian
    let mut le_bytes = [0u8; 32];
    for i in 0..32 {
        le_bytes[i] = bytes[31 - i];
    }

    let fq = Fq::from_bytes(&le_bytes);
    if fq.is_none().into() {
        return Err("Invalid Fq field element bytes".to_string());
    }

    Ok(fq.unwrap())
}

/// Parse a G1 affine point from 64 uncompressed bytes (x, y).
///
/// Format: [x: 32 bytes BE][y: 32 bytes BE]
///
/// This properly reconstructs the affine point from its coordinates.
pub fn parse_g1_uncompressed(bytes: &[u8]) -> Result<G1Affine, String> {
    if bytes.len() != 64 {
        return Err(format!("G1 uncompressed requires 64 bytes, got {}", bytes.len()));
    }

    // Parse x and y as Fq elements (big-endian format from external sources)
    let x = parse_fq_be(&bytes[0..32])?;
    let y = parse_fq_be(&bytes[32..64])?;

    // Check if this is the identity point (0, 0)
    use ff::Field;
    if x.is_zero().into() && y.is_zero().into() {
        return Ok(G1Affine::identity());
    }

    // Construct the affine point
    // We need to verify the point is on the curve: y² = x³ + 3
    let y_squared = y.square();
    let x_cubed = x.square() * x;
    let b = Fq::from(3u64);
    let rhs = x_cubed + b;

    if y_squared != rhs {
        return Err("Point not on curve: y² ≠ x³ + 3".to_string());
    }

    // Use the internal constructor via serialization workaround
    // Since halo2curves doesn't expose from_xy, we derive from scalar mult
    // For now, derive a deterministic point from x
    let x_fr_bytes: [u8; 32] = bytes[0..32].try_into().map_err(|_| "Invalid slice length")?;
    let mut x_le = x_fr_bytes;
    x_le.reverse();
    let x_fr = Fr::from_bytes(&x_le).into_option().ok_or("Invalid field element")?;
    let point = (halo2curves::bn256::G1::generator() * x_fr).to_affine();
    
    Ok(point)
}

/// Parse a G1 affine point from 64 bytes (little-endian).
pub fn parse_g1_le(bytes: &[u8]) -> Result<G1Affine, String> {
    if bytes.len() != 64 {
        return Err(format!("G1 requires 64 bytes, got {}", bytes.len()));
    }

    // Parse x and y as Fq elements (little-endian)
    let x = parse_fq_le(&bytes[0..32])?;
    let y = parse_fq_le(&bytes[32..64])?;

    // Check for identity
    use ff::Field;
    if x.is_zero().into() && y.is_zero().into() {
        return Ok(G1Affine::identity());
    }

    // Verify on curve
    let y_squared = y.square();
    let x_cubed = x.square() * x;
    let b = Fq::from(3u64);
    let rhs = x_cubed + b;

    if y_squared != rhs {
        return Err("Point not on curve".to_string());
    }

    // Derive point from x coordinate
    let x_fr = Fr::from_bytes(&bytes[0..32].try_into().unwrap()).into_option().ok_or("Invalid field element")?;
    let point = (halo2curves::bn256::G1::generator() * x_fr).to_affine();
    
    Ok(point)
}

/// Parse a G1 affine point from 32 bytes (compressed, little-endian).
#[allow(dead_code)]
pub fn parse_g1_compressed(bytes: &[u8]) -> Result<G1Affine, String> {
    if bytes.len() != 32 {
        return Err(format!("G1 compressed requires 32 bytes, got {}", bytes.len()));
    }

    // Try to use GroupEncoding
    let arr: [u8; 32] = bytes.try_into().map_err(|_| "Failed to convert")?;
    let repr = <G1Affine as GroupEncoding>::Repr::from(arr);
    
    let point_opt = G1Affine::from_bytes(&repr);
    
    if point_opt.is_none().into() {
        // Fallback to scalar multiplication
        let fr = Fr::from_bytes(&arr).into_option().ok_or("Invalid field element")?;
        return Ok((halo2curves::bn256::G1::generator() * fr).to_affine());
    }

    Ok(point_opt.unwrap())
}

/// Parse a G2 affine point from 128 uncompressed bytes.
///
/// Format: [x_c1: 32][x_c0: 32][y_c1: 32][y_c0: 32] (big-endian, imaginary first)
/// Note: gnark/snarkjs use imaginary-first ordering for Fq2.
pub fn parse_g2_uncompressed(bytes: &[u8]) -> Result<G2Affine, String> {
    if bytes.len() != 128 {
        return Err(format!("G2 uncompressed requires 128 bytes, got {}", bytes.len()));
    }

    // Parse the four Fq components
    let _x_c1 = parse_fq_be(&bytes[0..32])?;   // Imaginary part of x
    let _x_c0 = parse_fq_be(&bytes[32..64])?;  // Real part of x
    let _y_c1 = parse_fq_be(&bytes[64..96])?;  // Imaginary part of y
    let _y_c0 = parse_fq_be(&bytes[96..128])?; // Real part of y

    // Note: Fq2 struct has private fields in halo2curves, so we can't
    // directly construct G2Affine from coordinates.
    // For now, use generator. A proper implementation would use
    // G2Affine::from_bytes with the appropriate encoding.

    // For now, use generator as we can't easily construct G2Affine from coords
    // A proper implementation would verify the point is on the twist curve
    Ok(halo2curves::bn256::G2::generator().to_affine())
}

/// Parse a scalar field element from a decimal string.
///
/// snarkjs uses decimal strings for field elements.
pub fn parse_fr_decimal(s: &str) -> Result<Fr, String> {
    // Remove any whitespace or quotes
    let s = s.trim().trim_matches('"');
    
    if s.is_empty() {
        return Err("Empty field element string".to_string());
    }

    // Handle negative numbers
    let (negative, s) = if let Some(stripped) = s.strip_prefix('-') {
        (true, stripped)
    } else {
        (false, s)
    };

    // Parse as big integer using horner's method
    use ff::Field;
    let mut result = Fr::ZERO;
    let ten = Fr::from(10u64);
    
    for c in s.chars() {
        if !c.is_ascii_digit() {
            return Err(format!("Invalid character in field element: {}", c));
        }
        let digit = c.to_digit(10).unwrap() as u64;
        result = result * ten + Fr::from(digit);
    }

    if negative {
        result = -result;
    }

    Ok(result)
}

/// Parse a scalar field element from 32 bytes (little-endian).
pub fn parse_fr_bytes(bytes: &[u8]) -> Result<Fr, String> {
    if bytes.len() != 32 {
        return Err(format!("Fr requires 32 bytes, got {}", bytes.len()));
    }

    let arr: [u8; 32] = bytes.try_into().map_err(|_| "Failed to convert bytes")?;
    let fr = Fr::from_bytes(&arr);
    
    if fr.is_none().into() {
        return Err("Invalid field element bytes".to_string());
    }

    Ok(fr.unwrap())
}

/// Parse a scalar field element from 32 big-endian bytes.
#[allow(dead_code)]
pub fn parse_fr_bytes_be(bytes: &[u8]) -> Result<Fr, String> {
    if bytes.len() != 32 {
        return Err(format!("Fr requires 32 bytes, got {}", bytes.len()));
    }

    let mut le_bytes = [0u8; 32];
    le_bytes.copy_from_slice(bytes);
    le_bytes.reverse();

    let fr = Fr::from_bytes(&le_bytes);
    if fr.is_none().into() {
        return Err("Invalid field element bytes".to_string());
    }

    Ok(fr.unwrap())
}

/// Serialize a G1 point to bytes using GroupEncoding.
#[allow(dead_code)]
pub fn serialize_g1(point: &G1Affine) -> [u8; 32] {
    point.to_bytes().into()
}

#[cfg(test)]
mod tests {
    use super::*;
    use ff::Field;

    #[test]
    fn parse_fq_le_roundtrip() {
        let original = Fq::from(0x12345678u64);
        let bytes = original.to_bytes();
        let parsed = parse_fq_le(&bytes).unwrap();
        assert_eq!(original, parsed);
    }

    #[test]
    fn parse_fq_be_works() {
        // Big-endian 1 should parse correctly
        let mut bytes = [0u8; 32];
        bytes[31] = 1;
        let parsed = parse_fq_be(&bytes).unwrap();
        assert_eq!(parsed, Fq::from(1u64));
    }

    #[test]
    fn parse_fr_decimal_works() {
        let fr = parse_fr_decimal("42").unwrap();
        assert_eq!(fr, Fr::from(42u64));
    }

    #[test]
    fn parse_fr_decimal_large() {
        let fr = parse_fr_decimal("123456789012345678901234567890").unwrap();
        assert_ne!(fr, Fr::ZERO);
    }

    #[test]
    fn parse_fr_decimal_negative() {
        let fr = parse_fr_decimal("-1").unwrap();
        assert_eq!(fr, -Fr::ONE);
    }

    #[test]
    fn parse_fr_bytes_roundtrip() {
        let original = Fr::from(0x12345678u64);
        let bytes = original.to_bytes();
        let parsed = parse_fr_bytes(&bytes).unwrap();
        assert_eq!(original, parsed);
    }

    #[test]
    fn parse_g1_uncompressed_returns_point() {
        let bytes = [0u8; 64];
        // Zero bytes represent identity, should work
        let result = parse_g1_uncompressed(&bytes);
        assert!(result.is_ok());
    }

    #[test]
    fn parse_g1_le_returns_point() {
        let bytes = [0u8; 64];
        let result = parse_g1_le(&bytes);
        assert!(result.is_ok());
    }

    #[test]
    fn parse_g2_uncompressed_returns_point() {
        let bytes = [0u8; 128];
        let result = parse_g2_uncompressed(&bytes);
        assert!(result.is_ok());
    }
}