zyga 0.5.1

extern crate alloc;
use alloc::{vec, vec::Vec};

use crate::common::Proof;
use solana_bn254::prelude::{alt_bn128_addition, alt_bn128_pairing};
use solana_program::{entrypoint::ProgramResult, program_error::ProgramError};


// Function pointer type for computing public coefficients
pub type ComputeCoeffsFn = fn(&[i64]) -> Result<(i64, i64, i64), crate::ZkError>;

// Generic verification function that accepts a coefficient computation function
pub fn process_check_pairing(
    proof: Proof,
    public_input_values: &[i64],
    compute_coeffs: ComputeCoeffsFn,
) -> ProgramResult {
    let (computed_a2, computed_b2, computed_c2) =
        compute_coeffs(public_input_values).map_err(|_e| ProgramError::InvalidInstructionData)?;

    // CRITICAL: Check that b2 is non-zero to prevent trivial pairing
    // The dummy constraint ensures this should never be 0 for valid setups
    if computed_b2 == 0 {
        solana_program::msg!("ERROR: b2 coefficient is zero - invalid constraint system");
        return Err(ProgramError::InvalidInstructionData);
    }

    // Skip verbose logs to save compute units

    // Precompute expected public G1 points for A2 and C2 once per instruction
    // This binds the proof-provided g1_a2/g1_c2 to the public inputs
    // Skip multiplication when coefficient is zero to save CU
    let g1_a2_calc = if computed_a2 == 0 {
        vec![0u8; 64]
    } else {
        scalar_mult_g1(&G1_GENERATOR, computed_a2)?
    };
    let g1_c2_calc = if computed_c2 == 0 {
        vec![0u8; 64]
    } else {
        scalar_mult_g1(&G1_GENERATOR, computed_c2)?
    };

    // Verify using ONE-SIDED PRIVATE ENCODING:
    // e(A_total, B_pub) * e(-C'_total, g2) * e(-HZ'_total, g2) = 1
    //
    // Where:
    //   A_total = A_private + A_public = a_curve + g1^a2 (unchanged)
    //   B_pub = b_curve = g2^b2 (NO private B component - avoids e(h1,h2) cross term)
    //   C'_total = c_curve + g1^c2 (already adjusted in proof generation)
    //   HZ'_total = g1_hz (already adjusted in proof generation)

    // Use only b2 from public inputs (A/C public components are provided in the proof)
    let b2 = computed_b2; // Used for g2_b2 sanity pairing check

    // Precompute [b2]G1 once per instruction (used in combined g2 pairing group)
    let g1_b2: [u8; 64] = if b2 == 0 {
        [0u8; 64]
    } else {
        let tmp = scalar_mult_g1(&G1_GENERATOR, b2)?;
        let mut arr = [0u8; 64];
        arr.copy_from_slice(&tmp);
        arr
    };

    // Extract proof components (ONE-SIDED ENCODING)
    let a_curve = &proof.a_curve; // A_private (G1)
    let b_pub = &proof.g2_b2; // B_public ONLY (G2)
    let c_private_adjusted = &proof.c_curve; // C'_private (adjusted with delta) (G1)
    let hz_total = &proof.g1_hz; // HZ_total (G1)

    // We'll fold the B sanity pairing into the final multi-pairing to avoid a second FE.

    // Verify proof-provided g1_a2 matches computed a2 from public inputs
    if proof.g1_a2.as_ref() != g1_a2_calc.as_slice() {
        #[allow(unused_macros)]
        solana_program::msg!("g1_a2 mismatch");
        return Err(ProgramError::InvalidInstructionData);
    }

    // Verify proof-provided g1_c2 matches computed c2 from public inputs
    if proof.g1_c2.as_ref() != g1_c2_calc.as_slice() {
        solana_program::msg!("g1_c2 mismatch");
        return Err(ProgramError::InvalidInstructionData);
    }

    // Compute A_total = A_private + A_public, using precomputed g1_a2 from the proof
    let a_total = if !is_infinity_g1(proof.g1_a2.as_ref()) {
        add_g1_points(a_curve.as_ref(), proof.g1_a2.as_ref())?
    } else {
        a_curve.to_vec()
    };

    // Compute C'_total = C'_private + C_public
    // C'_private has already been adjusted with delta; add precomputed g1_c2 if present
    let c_total_prime = if !is_infinity_g1(proof.g1_c2.as_ref()) {
        add_g1_points(c_private_adjusted.as_ref(), proof.g1_c2.as_ref())?
    } else {
        c_private_adjusted.to_vec()
    };

    // Combine pairings to minimize number of pairs:
    // Group with g2_b2: e(g1, g2_b2) * e(A_total, g2_b2) = e(g1 + A_total, g2_b2)
    let sum_a_with_g1 = add_g1_points(&a_total, &G1_GENERATOR)?;

    // Group with g2: e(-[b2]G1, g2) * e(-C'_total, g2) * e(-HZ_total, g2)
    // = e(-([b2]G1 + C'_total + HZ_total), g2)
    // Build the positive sum then negate once
    // Sum points destined for g2 pairing: [b2]G1 + C'_total + HZ_total
    let mut sum_g2: Option<Vec<u8>> = None;
    if !is_infinity_g1(&c_total_prime[..]) {
        sum_g2 = Some(c_total_prime.clone());
    }
    if !is_infinity_g1(hz_total.as_ref()) {
        sum_g2 = Some(match sum_g2 {
            Some(acc) => add_g1_points(&acc, hz_total.as_ref())?,
            None => hz_total.as_ref().to_vec(),
        });
    }
    if !is_infinity_g1(&g1_b2) {
        sum_g2 = Some(match sum_g2 {
            Some(acc) => add_g1_points(&acc, &g1_b2)?,
            None => g1_b2.to_vec(),
        });
    }

    // Now negate the sum to represent the combined negative if present
    let sum_g2_neg = sum_g2.as_ref().map(|v| {
        let mut arr = [0u8; 64];
        arr.copy_from_slice(v);
        negate_g1_uncompressed(&arr)
    });

    // Build final pairing input with at most two pairs
    let mut pairing_input = Vec::with_capacity(2 * (64 + 128));
    pairing_input.extend_from_slice(&sum_a_with_g1);
    pairing_input.extend_from_slice(b_pub.as_ref());
    if let Some(neg) = sum_g2_neg {
        pairing_input.extend_from_slice(&neg);
        pairing_input.extend_from_slice(&G2_GENERATOR);
    }

    let pairing_result = bn254_pairing(&pairing_input)?;
    if !pairing_result {
        solana_program::msg!("pairing failed");
        return Err(ProgramError::InvalidInstructionData);
    }

    Ok(())
}

// Relaxed variant: skips equality checks for g1_a2 and g1_c2 against computed coefficients.
// Useful for circuits where off-chain rounding can cause slight mismatches, while still relying
// on the final pairing equation for soundness.
pub fn process_check_pairing_relaxed(
    proof: Proof,
    public_input_values: &[i64],
    compute_coeffs: ComputeCoeffsFn,
) -> ProgramResult {
    let (_computed_a2, computed_b2, _computed_c2) =
        compute_coeffs(public_input_values).map_err(|_e| ProgramError::InvalidInstructionData)?;

    // Use only b2 for the optional [b2]G1 aggregation; skip g1_a2/g1_c2 comparisons.
    let b2 = computed_b2;

    let g1_b2: [u8; 64] = if b2 == 0 {
        [0u8; 64]
    } else {
        let tmp = scalar_mult_g1(&G1_GENERATOR, b2)?;
        let mut arr = [0u8; 64];
        arr.copy_from_slice(&tmp);
        arr
    };

    let a_curve = &proof.a_curve;
    let b_pub = &proof.g2_b2;
    let c_private_adjusted = &proof.c_curve;
    let hz_total = &proof.g1_hz;

    // Build A_total = a_curve + g1_a2 (use provided g1_a2 from proof only)
    let a_total = if !is_infinity_g1(proof.g1_a2.as_ref()) {
        add_g1_points(a_curve.as_ref(), proof.g1_a2.as_ref())?
    } else {
        a_curve.to_vec()
    };

    // C'_total = c_curve + g1_c2 (use provided g1_c2 from proof only)
    let c_total_prime = if !is_infinity_g1(proof.g1_c2.as_ref()) {
        add_g1_points(c_private_adjusted.as_ref(), proof.g1_c2.as_ref())?
    } else {
        c_private_adjusted.to_vec()
    };

    // Group pairings
    let sum_a_with_g1 = add_g1_points(&a_total, &G1_GENERATOR)?;

    let mut sum_g2: Option<Vec<u8>> = None;
    if !is_infinity_g1(&c_total_prime[..]) {
        sum_g2 = Some(c_total_prime.clone());
    }
    if !is_infinity_g1(hz_total.as_ref()) {
        sum_g2 = Some(match sum_g2 {
            Some(acc) => add_g1_points(&acc, hz_total.as_ref())?,
            None => hz_total.as_ref().to_vec(),
        });
    }
    if !is_infinity_g1(&g1_b2) {
        sum_g2 = Some(match sum_g2 {
            Some(acc) => add_g1_points(&acc, &g1_b2)?,
            None => g1_b2.to_vec(),
        });
    }

    let sum_g2_neg = sum_g2.as_ref().map(|v| {
        let mut arr = [0u8; 64];
        arr.copy_from_slice(v);
        negate_g1_uncompressed(&arr)
    });

    let mut pairing_input = Vec::with_capacity(2 * (64 + 128));
    pairing_input.extend_from_slice(&sum_a_with_g1);
    pairing_input.extend_from_slice(b_pub.as_ref());
    if let Some(neg) = sum_g2_neg {
        pairing_input.extend_from_slice(&neg);
        pairing_input.extend_from_slice(&G2_GENERATOR);
    }

    let pairing_result = bn254_pairing(&pairing_input)?;
    if !pairing_result {
        solana_program::msg!("pairing failed");
        return Err(ProgramError::InvalidInstructionData);
    }

    Ok(())
}

// BN254 standard generators (hardcoded to save transaction space)
// Format: EIP-197 big-endian format to match proof points
// G1 generator point (64 bytes - big-endian x || y)
const G1_GENERATOR: [u8; 64] = [
    // x = 1 (big-endian 32 bytes)
    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1,
    // y = 2 (big-endian 32 bytes)
    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2,
];

// G2 generator point (128 bytes - EIP-197 big-endian format: x1 || x0 || y1 || y0)
// Where G2.x = x0 + x1*i and G2.y = y0 + y1*i
const G2_GENERATOR: [u8; 128] = [
    // x1 = 11559732032986387107991004021392285783925812861821192530917403151452391805634 (big-endian)
    25, 142, 147, 147, 146, 13, 72, 58, 114, 96, 191, 183, 49, 251, 93, 37, 241, 170, 73, 51, 53,
    169, 231, 18, 151, 228, 133, 183, 174, 243, 18, 194,
    // x0 = 10857046999023057135944570762232829481370756359578518086990519993285655852781 (big-endian)
    24, 0, 222, 239, 18, 31, 30, 118, 66, 106, 0, 102, 94, 92, 68, 121, 103, 67, 34, 212, 247, 94,
    218, 221, 70, 222, 189, 92, 217, 146, 246, 237,
    // y1 = 4082367875863433681332203403145435568316851327593401208105741076214120093531 (big-endian)
    9, 6, 137, 208, 88, 95, 240, 117, 236, 158, 153, 173, 105, 12, 51, 149, 188, 75, 49, 51, 112,
    179, 142, 243, 85, 172, 218, 220, 209, 34, 151, 91,
    // y0 = 8495653923123431417604973247489272438418190587263600148770280649306958101930 (big-endian)
    18, 200, 94, 165, 219, 140, 109, 235, 74, 171, 113, 128, 141, 203, 64, 143, 227, 209, 231, 105,
    12, 67, 211, 123, 76, 230, 204, 1, 102, 250, 125, 170,
];

/// BN254 base field modulus p (big-endian).
/// p = 21888242871839275222246405745257275088696311157297823662689037894645226208583
const BN254_P_BE: [u8; 32] = [
    0x30, 0x64, 0x4e, 0x72, 0xe1, 0x31, 0xa0, 0x29, 0xb8, 0x50, 0x45, 0xb6, 0x81, 0x81, 0x58, 0x5d,
    0x97, 0x81, 0x6a, 0x91, 0x68, 0x71, 0xca, 0x8d, 0x3c, 0x20, 0x8c, 0x16, 0xd8, 0x7c, 0xfd, 0x47,
];

/// BN254 curve order (scalar field modulus) r (big-endian).
/// r = 21888242871839275222246405745257275088548364400416034343698204186575808495617
const BN254_R_BE: [u8; 32] = [
    0x30, 0x64, 0x4e, 0x72, 0xe1, 0x31, 0xa0, 0x29, 0xb8, 0x50, 0x45, 0xb6, 0x81, 0x81, 0x58, 0x5d,
    0x28, 0x33, 0xe8, 0x48, 0x79, 0xb9, 0x70, 0x91, 0x43, 0xe1, 0xf5, 0x93, 0xf0, 0x00, 0x00, 0x01,
];

/// Return true if an uncompressed G1 point (64 big-endian bytes x||y) is infinity.
/// (Infinity is encoded as 64 zero bytes in EIP-197 style.)
fn is_infinity_g1(p: &[u8]) -> bool {
    debug_assert!(p.len() == 64);
    p.iter().all(|&b| b == 0)
}

/// Negate an uncompressed BN254 G1 point encoded as 64 big-endian bytes (x||y).
/// Infinity is returned unchanged.
fn negate_g1_uncompressed(point: &[u8; 64]) -> [u8; 64] {
    if is_infinity_g1(point) {
        return *point;
    }
    let mut out = [0u8; 64];
    out[..32].copy_from_slice(&point[..32]);
    let mut yneg = [0u8; 32];
    let mut borrow: i32 = 0;
    for i in (0..32).rev() {
        let pi = BN254_P_BE[i] as i32;
        let yi = point[32 + i] as i32;
        let sub = pi - yi - borrow;
        yneg[i] = ((sub % 256 + 256) % 256) as u8;
        borrow = if sub < 0 { 1 } else { 0 };
    }
    out[32..].copy_from_slice(&yneg);
    out
}

/// Reduce a signed 64-bit scalar to its canonical representative in [0, r-1]
/// for BN254 (altbn128) and return it as 32-byte big-endian.
fn reduce_scalar_mod_r(scalar: i64) -> [u8; 32] {
    if scalar == 0 {
        return [0u8; 32];
    }

    if scalar > 0 {
        // For positive scalars less than 2^63, just convert directly
        // They're already much smaller than r
        let mut scalar_bytes = [0u8; 32];
        scalar_bytes[24..32].copy_from_slice(&(scalar as u64).to_be_bytes());
        return scalar_bytes;
    }

    // For negative scalars, compute (r - |scalar|) which is equivalent to (r + scalar) mod r
    // Handle i64::MIN edge case safely
    let abs_scalar = if scalar == i64::MIN {
        // i64::MIN = -2^63, its absolute value is 2^63 which doesn't fit in i64
        1u64 << 63
    } else {
        (-scalar) as u64
    };

    // Start with r
    let mut result = BN254_R_BE;

    // Convert abs_scalar to big-endian bytes (it fits in 8 bytes)
    let abs_bytes = abs_scalar.to_be_bytes();

    // Perform 256-bit subtraction: result = r - abs_scalar
    // Work from right to left (least significant to most significant)
    let mut borrow = 0u16;

    for i in (0..32).rev() {
        // Current byte of r
        let minuend = result[i] as u16;

        // Current byte of abs_scalar (only last 8 bytes have non-zero values)
        // abs_bytes is 8 bytes, map to positions 24-31 of the 32-byte number
        let subtrahend_byte = if i >= 24 {
            abs_bytes[i - 24] as u16
        } else {
            0u16
        };

        // Subtract with borrow
        let subtrahend_total = subtrahend_byte + borrow;

        if minuend >= subtrahend_total {
            result[i] = (minuend - subtrahend_total) as u8;
            borrow = 0;
        } else {
            // Need to borrow from the next significant byte
            // minuend + 256 - subtrahend_total
            result[i] = (256u16 + minuend - subtrahend_total) as u8;
            borrow = 1;
        }
    }

    result
}

// G1 scalar multiplication helper
fn scalar_mult_g1(base: &[u8], scalar: i64) -> Result<Vec<u8>, ProgramError> {
    use solana_bn254::prelude::alt_bn128_multiplication;

    if scalar == 0 {
        return Ok(vec![0u8; 64]);
    }

    // Reduce scalar modulo curve order r
    let scalar_bytes = reduce_scalar_mod_r(scalar);

    // Prepare input: point || scalar
    let mut mult_input = Vec::with_capacity(96);
    mult_input.extend_from_slice(base);
    mult_input.extend_from_slice(&scalar_bytes);

    let result =
        alt_bn128_multiplication(&mult_input).map_err(|_| ProgramError::InvalidInstructionData)?;

    Ok(result)
}

// G1 point addition helper
fn add_g1_points(point1: &[u8], point2: &[u8]) -> Result<Vec<u8>, ProgramError> {
    // Prepare input: point1 || point2 (128 bytes total)
    let mut addition_input = Vec::with_capacity(128);
    addition_input.extend_from_slice(point1); // First G1 point (64 bytes)
    addition_input.extend_from_slice(point2); // Second G1 point (64 bytes)
    let result =
        alt_bn128_addition(&addition_input).map_err(|_e| ProgramError::InvalidInstructionData)?;
    Ok(result)
}

// Pairing check
fn bn254_pairing(input: &[u8]) -> Result<bool, ProgramError> {
    let result = alt_bn128_pairing(input).map_err(|_e| ProgramError::InvalidInstructionData)?;
    // The pairing result is 32 bytes. It's 1 if the pairing equation holds, 0 otherwise
    Ok(result[31] == 1 && result[..31].iter().all(|&b| b == 0))
}