//! ZK Proof generation and verification.
//!
//! This module provides the cryptographic proof system for ZK-rollups,
//! supporting multiple proving backends (Groth16, PLONK, etc.).
//!
//! # Proof Lifecycle
//!
//! ```text
//! Setup Phase (one-time):
//!   Circuit → Setup → (ProvingKey, VerificationKey)
//!
//! Proving Phase (per batch):
//!   (Circuit, Witness, ProvingKey) → Prover → Proof
//!
//! Verification Phase (on-chain):
//!   (Proof, PublicInputs, VerificationKey) → Verifier → Accept/Reject
//! ```

use ark_bn254::{Bn254, Fr as ScalarField};
use ark_ff::{BigInteger, PrimeField};
use ark_groth16::ProvingKey as Groth16ProvingKey;
use ark_serialize::{CanonicalDeserialize, CanonicalSerialize};
use serde::{Deserialize, Serialize};
use thiserror::Error;

use crate::circuit::{Circuit, CircuitError};

/// Proof system backend selection.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub enum ProofSystemBackend {
    /// Groth16 - smallest proofs (~200 bytes), requires per-circuit trusted setup
    Groth16,
    /// PLONK - universal setup, medium proofs (~500 bytes)
    Plonk,
    /// STARK - no trusted setup, large proofs (~50KB), fast prover
    Stark,
}

impl Default for ProofSystemBackend {
    fn default() -> Self {
        // Groth16 is the default for production (smallest proofs)
        ProofSystemBackend::Groth16
    }
}

/// Verification key for proof validation.
#[derive(Clone)]
pub struct VerificationKey {
    /// Backend type
    pub backend: ProofSystemBackend,
    /// Serialized key data
    pub data: Vec<u8>,
}

impl VerificationKey {
    /// Creates a verification key from Groth16 verification key.
    pub fn from_groth16(vk: &ark_groth16::VerifyingKey<Bn254>) -> Result<Self, ProofError> {
        let mut data = Vec::new();
        vk.serialize_compressed(&mut data)
            .map_err(|e| ProofError::SerializationError(e.to_string()))?;

        Ok(Self {
            backend: ProofSystemBackend::Groth16,
            data,
        })
    }

    /// Deserializes the Groth16 verification key.
    pub fn to_groth16(&self) -> Result<ark_groth16::VerifyingKey<Bn254>, ProofError> {
        if self.backend != ProofSystemBackend::Groth16 {
            return Err(ProofError::BackendMismatch);
        }

        ark_groth16::VerifyingKey::deserialize_compressed(&self.data[..])
            .map_err(|e| ProofError::DeserializationError(e.to_string()))
    }

    /// Returns the key size in bytes.
    pub fn size(&self) -> usize {
        self.data.len()
    }

    /// Serializes to bytes.
    pub fn to_bytes(&self) -> Vec<u8> {
        let mut result = vec![self.backend as u8];
        result.extend(&self.data);
        result
    }

    /// Deserializes from bytes.
    pub fn from_bytes(bytes: &[u8]) -> Result<Self, ProofError> {
        if bytes.is_empty() {
            return Err(ProofError::InvalidKeyFormat);
        }

        let backend = match bytes[0] {
            0 => ProofSystemBackend::Groth16,
            1 => ProofSystemBackend::Plonk,
            2 => ProofSystemBackend::Stark,
            _ => return Err(ProofError::InvalidKeyFormat),
        };

        Ok(Self {
            backend,
            data: bytes[1..].to_vec(),
        })
    }
}

impl std::fmt::Debug for VerificationKey {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("VerificationKey")
            .field("backend", &self.backend)
            .field("size", &self.data.len())
            .finish()
    }
}

/// Proving key for proof generation.
pub struct ProvingKey {
    /// Backend type
    pub backend: ProofSystemBackend,
    /// Serialized key data
    pub data: Vec<u8>,
}

impl ProvingKey {
    /// Creates a proving key from Groth16 proving key.
    pub fn from_groth16(pk: &Groth16ProvingKey<Bn254>) -> Result<Self, ProofError> {
        let mut data = Vec::new();
        pk.serialize_compressed(&mut data)
            .map_err(|e| ProofError::SerializationError(e.to_string()))?;

        Ok(Self {
            backend: ProofSystemBackend::Groth16,
            data,
        })
    }

    /// Deserializes the Groth16 proving key.
    pub fn to_groth16(&self) -> Result<Groth16ProvingKey<Bn254>, ProofError> {
        if self.backend != ProofSystemBackend::Groth16 {
            return Err(ProofError::BackendMismatch);
        }

        Groth16ProvingKey::deserialize_compressed(&self.data[..])
            .map_err(|e| ProofError::DeserializationError(e.to_string()))
    }

    /// Returns the key size in bytes.
    pub fn size(&self) -> usize {
        self.data.len()
    }
}

impl std::fmt::Debug for ProvingKey {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ProvingKey")
            .field("backend", &self.backend)
            .field("size", &self.data.len())
            .finish()
    }
}

/// Zero-knowledge proof.
#[derive(Clone, Serialize, Deserialize)]
pub struct Proof {
    /// Backend used to generate this proof
    pub backend: ProofSystemBackend,
    /// Serialized proof data
    pub data: Vec<u8>,
    /// Public inputs (for verification)
    pub public_inputs: Vec<[u8; 32]>,
}

impl Proof {
    /// Creates a proof from Groth16 proof.
    pub fn from_groth16(
        proof: &ark_groth16::Proof<Bn254>,
        public_inputs: Vec<ScalarField>,
    ) -> Result<Self, ProofError> {
        let mut data = Vec::new();
        proof
            .serialize_compressed(&mut data)
            .map_err(|e| ProofError::SerializationError(e.to_string()))?;

        // Convert field elements to bytes
        let public_inputs_bytes: Vec<[u8; 32]> = public_inputs
            .iter()
            .map(|f| {
                let mut bytes = [0u8; 32];
                let f_bytes = f.into_bigint().to_bytes_le();
                bytes[..f_bytes.len().min(32)].copy_from_slice(&f_bytes[..f_bytes.len().min(32)]);
                bytes
            })
            .collect();

        Ok(Self {
            backend: ProofSystemBackend::Groth16,
            data,
            public_inputs: public_inputs_bytes,
        })
    }

    /// Deserializes the Groth16 proof.
    pub fn to_groth16(&self) -> Result<ark_groth16::Proof<Bn254>, ProofError> {
        if self.backend != ProofSystemBackend::Groth16 {
            return Err(ProofError::BackendMismatch);
        }

        ark_groth16::Proof::deserialize_compressed(&self.data[..])
            .map_err(|e| ProofError::DeserializationError(e.to_string()))
    }

    /// Returns the proof size in bytes.
    pub fn size(&self) -> usize {
        self.data.len()
    }

    /// Serializes to bytes.
    pub fn to_bytes(&self) -> Vec<u8> {
        let mut result = vec![self.backend as u8];

        // Write data length and data
        let data_len = (self.data.len() as u32).to_le_bytes();
        result.extend(&data_len);
        result.extend(&self.data);

        // Write public inputs
        let input_count = (self.public_inputs.len() as u32).to_le_bytes();
        result.extend(&input_count);
        for input in &self.public_inputs {
            result.extend(input);
        }

        result
    }

    /// Deserializes from bytes.
    pub fn from_bytes(bytes: &[u8]) -> Result<Self, ProofError> {
        if bytes.len() < 9 {
            return Err(ProofError::InvalidProofFormat);
        }

        let backend = match bytes[0] {
            0 => ProofSystemBackend::Groth16,
            1 => ProofSystemBackend::Plonk,
            2 => ProofSystemBackend::Stark,
            _ => return Err(ProofError::InvalidProofFormat),
        };

        let data_len = u32::from_le_bytes([bytes[1], bytes[2], bytes[3], bytes[4]]) as usize;
        if bytes.len() < 5 + data_len + 4 {
            return Err(ProofError::InvalidProofFormat);
        }

        let data = bytes[5..5 + data_len].to_vec();

        let input_offset = 5 + data_len;
        let input_count =
            u32::from_le_bytes([bytes[input_offset], bytes[input_offset + 1], bytes[input_offset + 2], bytes[input_offset + 3]])
                as usize;

        let mut public_inputs = Vec::with_capacity(input_count);
        let mut offset = input_offset + 4;
        for _ in 0..input_count {
            if bytes.len() < offset + 32 {
                return Err(ProofError::InvalidProofFormat);
            }
            let mut input = [0u8; 32];
            input.copy_from_slice(&bytes[offset..offset + 32]);
            public_inputs.push(input);
            offset += 32;
        }

        Ok(Self {
            backend,
            data,
            public_inputs,
        })
    }
}

impl std::fmt::Debug for Proof {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Proof")
            .field("backend", &self.backend)
            .field("size", &self.data.len())
            .field("public_inputs", &self.public_inputs.len())
            .finish()
    }
}

/// Proof system for generating and verifying proofs.
pub struct ProofSystem {
    backend: ProofSystemBackend,
}

impl ProofSystem {
    /// Creates a new proof system with the specified backend.
    pub fn new(backend: ProofSystemBackend) -> Self {
        Self { backend }
    }

    /// Creates a proof system with the default backend (Groth16).
    pub fn groth16() -> Self {
        Self::new(ProofSystemBackend::Groth16)
    }

    /// Returns the backend used.
    pub fn backend(&self) -> ProofSystemBackend {
        self.backend
    }

    /// Generates proving and verification keys for a circuit.
    ///
    /// This is the "trusted setup" phase for Groth16.
    /// For production, use a multi-party computation ceremony.
    pub fn setup<C: Circuit>(&self, _circuit: &C) -> Result<(ProvingKey, VerificationKey), ProofError> {
        match self.backend {
            ProofSystemBackend::Groth16 => {
                // In a real implementation, we would:
                // 1. Convert our Circuit trait to ark_relations::r1cs::ConstraintSynthesizer
                // 2. Run Groth16::circuit_specific_setup

                // For now, return placeholder keys
                // Real implementation would use:
                // let (pk, vk) = Groth16::<Bn254>::circuit_specific_setup(circuit, rng)?;

                Ok((
                    ProvingKey {
                        backend: self.backend,
                        data: vec![0u8; 1024], // Placeholder
                    },
                    VerificationKey {
                        backend: self.backend,
                        data: vec![0u8; 256], // Placeholder
                    },
                ))
            }
            ProofSystemBackend::Plonk => {
                Err(ProofError::UnsupportedBackend("PLONK not yet implemented".into()))
            }
            ProofSystemBackend::Stark => {
                Err(ProofError::UnsupportedBackend("STARK not yet implemented".into()))
            }
        }
    }

    /// Generates a proof for a circuit with the given proving key.
    pub fn prove<C: Circuit>(
        &self,
        _circuit: &C,
        _proving_key: &ProvingKey,
    ) -> Result<Proof, ProofError> {
        match self.backend {
            ProofSystemBackend::Groth16 => {
                // In a real implementation:
                // let proof = Groth16::<Bn254>::prove(&pk, circuit, rng)?;

                Ok(Proof {
                    backend: self.backend,
                    data: vec![0u8; 192], // Groth16 proof is ~192 bytes
                    public_inputs: Vec::new(),
                })
            }
            _ => Err(ProofError::UnsupportedBackend(format!("{:?}", self.backend))),
        }
    }

    /// Verifies a proof against a verification key.
    pub fn verify(
        &self,
        _proof: &Proof,
        _verification_key: &VerificationKey,
    ) -> Result<bool, ProofError> {
        match self.backend {
            ProofSystemBackend::Groth16 => {
                // In a real implementation:
                // let pvk = PreparedVerifyingKey::from(&vk);
                // let result = Groth16::<Bn254>::verify_with_processed_vk(&pvk, &public_inputs, &proof)?;

                // Placeholder: always return true for testing
                Ok(true)
            }
            _ => Err(ProofError::UnsupportedBackend(format!("{:?}", self.backend))),
        }
    }

    /// Estimates proof generation time for a circuit.
    pub fn estimate_proving_time<C: Circuit>(circuit: &C) -> std::time::Duration {
        // Rough estimates based on circuit complexity
        let config = circuit.config();
        let base_ms = 100; // Base overhead
        let per_constraint_us = 10; // Microseconds per constraint

        // Estimate constraints based on batch size
        let estimated_constraints = config.max_batch_size * 1000;

        std::time::Duration::from_millis(
            base_ms + (estimated_constraints as u64 * per_constraint_us / 1000),
        )
    }

    /// Estimates proof verification time.
    pub fn estimate_verification_time(backend: ProofSystemBackend) -> std::time::Duration {
        match backend {
            ProofSystemBackend::Groth16 => std::time::Duration::from_millis(10),
            ProofSystemBackend::Plonk => std::time::Duration::from_millis(15),
            ProofSystemBackend::Stark => std::time::Duration::from_millis(30),
        }
    }

    /// Returns the expected proof size for a backend.
    pub fn proof_size(backend: ProofSystemBackend) -> usize {
        match backend {
            ProofSystemBackend::Groth16 => 192,      // ~200 bytes
            ProofSystemBackend::Plonk => 512,        // ~500 bytes
            ProofSystemBackend::Stark => 50 * 1024, // ~50KB
        }
    }
}

impl Default for ProofSystem {
    fn default() -> Self {
        Self::groth16()
    }
}

/// Proof system errors.
#[derive(Debug, Error)]
pub enum ProofError {
    #[error("Proof generation failed: {0}")]
    ProvingError(String),

    #[error("Proof verification failed: {0}")]
    VerificationError(String),

    #[error("Setup failed: {0}")]
    SetupError(String),

    #[error("Serialization error: {0}")]
    SerializationError(String),

    #[error("Deserialization error: {0}")]
    DeserializationError(String),

    #[error("Backend mismatch")]
    BackendMismatch,

    #[error("Invalid proof format")]
    InvalidProofFormat,

    #[error("Invalid key format")]
    InvalidKeyFormat,

    #[error("Unsupported backend: {0}")]
    UnsupportedBackend(String),

    #[error("Circuit error: {0}")]
    CircuitError(#[from] CircuitError),
}

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

    #[test]
    fn test_proof_system_creation() {
        let ps = ProofSystem::groth16();
        assert_eq!(ps.backend(), ProofSystemBackend::Groth16);
    }

    #[test]
    fn test_proof_sizes() {
        assert!(ProofSystem::proof_size(ProofSystemBackend::Groth16) < 300);
        assert!(ProofSystem::proof_size(ProofSystemBackend::Plonk) < 1000);
        assert!(ProofSystem::proof_size(ProofSystemBackend::Stark) > 10000);
    }

    #[test]
    fn test_verification_time_estimates() {
        let groth16_time = ProofSystem::estimate_verification_time(ProofSystemBackend::Groth16);
        let plonk_time = ProofSystem::estimate_verification_time(ProofSystemBackend::Plonk);
        let stark_time = ProofSystem::estimate_verification_time(ProofSystemBackend::Stark);

        assert!(groth16_time < plonk_time);
        assert!(plonk_time < stark_time);
    }

    #[test]
    fn test_proof_serialization() {
        let proof = Proof {
            backend: ProofSystemBackend::Groth16,
            data: vec![1, 2, 3, 4, 5],
            public_inputs: vec![[0xab; 32], [0xcd; 32]],
        };

        let bytes = proof.to_bytes();
        let decoded = Proof::from_bytes(&bytes).unwrap();

        assert_eq!(decoded.backend, proof.backend);
        assert_eq!(decoded.data, proof.data);
        assert_eq!(decoded.public_inputs, proof.public_inputs);
    }

    #[test]
    fn test_verification_key_serialization() {
        let vk = VerificationKey {
            backend: ProofSystemBackend::Groth16,
            data: vec![1, 2, 3, 4, 5],
        };

        let bytes = vk.to_bytes();
        let decoded = VerificationKey::from_bytes(&bytes).unwrap();

        assert_eq!(decoded.backend, vk.backend);
        assert_eq!(decoded.data, vk.data);
    }
}
