zyga 0.5.1

use crate::common::{G1Point, G2Point, TrustedSetup};
use crate::dag::{ExprId, ExpressionDAG};
use crate::errors::ZkError;
use serde::{Deserialize, Serialize};
use std::collections::HashSet;

/// Proving key containing all static elements needed for proof generation
/// This is saved to a .zyga file for off-chain use
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ProvingKey {
    /// Constraint matrices (integer coefficients)
    pub a_matrix: Vec<Vec<f64>>,  // Keep as f64 for now to avoid breaking changes
    pub b_matrix: Vec<Vec<f64>>,  // These come from compilation_result which uses f64
    pub c_matrix: Vec<Vec<f64>>,  // TODO: Should be i64 in the future

    /// Pre-computed Lagrange evaluations at the evaluation point (MUST be f64)
    pub lagrange_a: Vec<f64>,
    pub lagrange_b: Vec<f64>,
    pub lagrange_c: Vec<f64>,

    /// Witness structure
    pub witness_dag: ExpressionDAG,
    pub witness_ids: Vec<ExprId>,
    pub witness_names: Vec<String>,

    /// Public variables
    pub public_variables: HashSet<String>,

    /// Trusted setup parameters
    pub trusted_setup: TrustedSetup,

    /// Metadata
    pub num_constraints: usize,
    pub num_variables: usize,
    pub evaluation_point: f64,  // This needs to be f64 for polynomial evaluation

    /// Environment dictionary for public values (integers)
    pub env_dict: std::collections::HashMap<String, i64>,
}

/// Verification key containing static elements needed for on-chain verification
/// This is exported to a proving_key.rs file
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct VerificationKey {
    /// Standard generators
    pub g1: G1Point,
    pub g2: G2Point,

    /// Trusted setup points
    pub h1: G1Point,  // g1^alpha
    pub h2: G2Point,  // g2^beta

    /// Number of constraints
    pub num_constraints: usize,
}

impl ProvingKey {
    /// Save proving key to a binary file
    pub fn save_to_file(&self, path: &std::path::Path) -> Result<(), ZkError> {
        let encoded = bincode::serialize(self)
            .map_err(|e| ZkError::SerializationError(format!("Failed to serialize proving key: {}", e)))?;

        std::fs::write(path, encoded)
            .map_err(|e| ZkError::IoError(format!("Failed to write proving key file: {}", e)))?;

        Ok(())
    }

    /// Load proving key from a binary file
    pub fn load_from_file(path: &std::path::Path) -> Result<Self, ZkError> {
        let data = std::fs::read(path)
            .map_err(|e| ZkError::IoError(format!("Failed to read proving key file: {}", e)))?;

        let proving_key = bincode::deserialize(&data)
            .map_err(|e| ZkError::SerializationError(format!("Failed to deserialize proving key: {}", e)))?;

        Ok(proving_key)
    }
}

impl VerificationKey {
    /// Generate Rust code for on-chain verification
    pub fn to_rust_code(&self, prefix: Option<&str>) -> String {
        let prefix_upper = prefix.map(|p| p.to_uppercase()).unwrap_or_default();
        let prefix_lower = prefix.map(|p| p.to_lowercase()).unwrap_or_default();
        let prefix_const = if prefix.is_some() {
            format!("{}_", prefix_upper)
        } else {
            String::new()
        };
        let prefix_fn = if prefix.is_some() {
            format!("{}_", prefix_lower)
        } else {
            String::new()
        };

        let mut code = String::new();

        // Header
        code.push_str("// Auto-generated proving key for on-chain verification\n");
        if let Some(p) = prefix {
            code.push_str(&format!("// Prefix: {}\n", p));
        }
        code.push_str("// Generated from constraint compilation\n");
        code.push_str("// DO NOT EDIT MANUALLY\n");
        code.push_str("\n");
        code.push_str("#![allow(dead_code)]\n");
        code.push_str("\n");
        code.push_str("use zyga::{G1Point, G2Point};\n");
        code.push_str("\n");

        // Number of constraints
        code.push_str(&format!("/// Number of constraints in the system\n"));
        code.push_str(&format!("pub const {}NUM_CONSTRAINTS: usize = {};\n", prefix_const, self.num_constraints));
        code.push_str("\n");

        // G1 generator
        code.push_str(&format!("/// G1 generator point\n"));
        code.push_str(&format!("pub const {}G1_GENERATOR: G1Point = G1Point(\n", prefix_const));
        code.push_str(&format!("    {:?}\n", self.g1.to_bytes()));
        code.push_str(");\n");
        code.push_str("\n");

        // G2 generator
        code.push_str(&format!("/// G2 generator point\n"));
        code.push_str(&format!("pub const {}G2_GENERATOR: G2Point = G2Point(\n", prefix_const));
        code.push_str(&format!("    {:?}\n", self.g2.to_bytes()));
        code.push_str(");\n");
        code.push_str("\n");

        // H1 (g1^alpha)
        code.push_str(&format!("/// Trusted setup h1 = g1^alpha\n"));
        code.push_str(&format!("pub const {}H1: G1Point = G1Point(\n", prefix_const));
        code.push_str(&format!("    {:?}\n", self.h1.to_bytes()));
        code.push_str(");\n");
        code.push_str("\n");

        // H2 (g2^beta)
        code.push_str(&format!("/// Trusted setup h2 = g2^beta\n"));
        code.push_str(&format!("pub const {}H2: G2Point = G2Point(\n", prefix_const));
        code.push_str(&format!("    {:?}\n", self.h2.to_bytes()));
        code.push_str(");\n");
        code.push_str("\n");

        // Helper function
        code.push_str(&format!("/// Get proving key elements\n"));
        code.push_str(&format!("pub fn get_{}proving_key() -> (G1Point, G2Point, G1Point, G2Point, usize) {{\n", prefix_fn));
        code.push_str(&format!("    (\n"));
        code.push_str(&format!("        {}G1_GENERATOR,\n", prefix_const));
        code.push_str(&format!("        {}G2_GENERATOR,\n", prefix_const));
        code.push_str(&format!("        {}H1,\n", prefix_const));
        code.push_str(&format!("        {}H2,\n", prefix_const));
        code.push_str(&format!("        {}NUM_CONSTRAINTS,\n", prefix_const));
        code.push_str(&format!("    )\n"));
        code.push_str("}\n");

        code
    }
}