# Copyright (c) 2025 Aaryan Guglani and FedZK Contributors
# SPDX-License-Identifier: MIT

"""
Zero-Knowledge Proof Verifier for FedZK.

This module contains the ZKVerifier class which handles verification of zero-knowledge
proofs for gradient updates in federated learning.
"""

import json
import os
import subprocess
import tempfile
from pathlib import Path
from typing import Any, Dict, List, Optional


class ZKVerifier:
    """
    Verifies zero-knowledge proofs for model updates in federated learning.
    
    This class can verify both dummy proofs (for testing) and real zero-knowledge
    proofs generated by snarkjs and Circom circuits.
    """

    def __init__(self, verification_key_path: str):
        """
        Initialize with path to verification key.
        
        Args:
            verification_key_path: Path to the verification key for the ZK circuit
        """
        self.verification_key_path = verification_key_path

    def verify_proof(self, proof: str, public_signals: List[Any]) -> bool:
        """
        Simulate verifying a ZK proof.
        
        For now, return True if the proof string starts with "dummy_proof_" and len(public_signals) > 0.
        In production, this would use snarkjs or a native verifier.
        
        Args:
            proof: The proof string generated by ZKProver
            public_signals: List of public signals/inputs for the proof
            
        Returns:
            Boolean indicating whether the proof is valid
        """
        # Simple validation for dummy implementation
        is_valid_proof_format = isinstance(proof, str) and proof.startswith("dummy_proof_")
        has_public_signals = isinstance(public_signals, list) and len(public_signals) > 0

        # In a real implementation, this would call snarkjs or a native verifier library
        # to verify the proof cryptographically against the verification key

        return is_valid_proof_format and has_public_signals

    def verify_real_proof(self, proof: Dict[str, Any], public_inputs: List[Any]) -> bool:
        """
        Verify a real zero-knowledge proof using snarkjs.
        """
        # Create a temporary directory for verification
        with tempfile.TemporaryDirectory() as temp_dir:
            temp_path = Path(temp_dir)
            proof_path = temp_path / "proof.json"
            public_path = temp_path / "public.json"
            with open(proof_path, "w") as f:
                json.dump(proof, f)
            with open(public_path, "w") as f:
                json.dump(public_inputs, f)
            try:
                result = subprocess.run([
                    "snarkjs", "groth16", "verify", self.verification_key_path,
                    str(public_path), str(proof_path)
                ], capture_output=True, text=True, check=False)
                if result.returncode == 0 and "OK" in result.stdout:
                    return True
                else:
                    return False
            except subprocess.SubprocessError:
                return False

    def verify_real_proof_secure(self, proof: Dict[str, Any], public_inputs: List[Any], verification_key_path: Optional[str] = None) -> bool:
        """
        Verify a secure zero-knowledge proof using snarkjs with the secure circuit.
        """
        if verification_key_path is None:
            verification_key_path = os.path.join(
                os.path.dirname(self.verification_key_path),
                "verification_key_secure.json"
            )
        if not os.path.exists(verification_key_path):
            raise FileNotFoundError(f"Secure verification key not found: {verification_key_path}")
        with tempfile.TemporaryDirectory() as temp_dir:
            temp_path = Path(temp_dir)
            proof_path = temp_path / "proof.json"
            public_path = temp_path / "public.json"
            with open(proof_path, "w") as f:
                json.dump(proof, f)
            with open(public_path, "w") as f:
                json.dump(public_inputs, f)
            try:
                result = subprocess.run([
                    "snarkjs", "groth16", "verify", verification_key_path,
                    str(public_path), str(proof_path)
                ], capture_output=True, text=True, check=False)
                if result.returncode == 0 and "OK" in result.stdout:
                    return True
                else:
                    return False
            except subprocess.SubprocessError:
                return False
