import Foundation
import SynorZk

/// Synor ZK SDK Examples for Swift
///
/// Demonstrates zero-knowledge proof operations:
/// - Circuit compilation (Circom)
/// - Proof generation and verification
/// - Multiple proving systems (Groth16, PLONK, STARK)
/// - Recursive proof composition
/// - Trusted setup ceremonies

@main
struct ZkExample {
    static func main() async throws {
        // Initialize client
        let config = ZkConfig(
            apiKey: ProcessInfo.processInfo.environment["SYNOR_API_KEY"] ?? "your-api-key",
            endpoint: "https://zk.synor.io/v1",
            timeout: 120, // ZK operations can be slow
            retries: 3,
            debug: false,
            defaultProvingSystem: .groth16,
            proveTimeout: 300,
            verifyTimeout: 30
        )

        let zk = SynorZk(config: config)

        do {
            // Check service health
            let healthy = try await zk.healthCheck()
            print("Service healthy: \(healthy)\n")

            // Run examples
            try await circuitExample(zk: zk)
            try await groth16Example(zk: zk)
            try await plonkExample(zk: zk)
            try await starkExample(zk: zk)
            try await recursiveExample(zk: zk)
            try await ceremonyExample(zk: zk)
        } catch {
            print("Error: \(error)")
        }

        await zk.close()
    }

    static func circuitExample(zk: SynorZk) async throws {
        print("=== Circuit Compilation ===")

        // Simple Circom circuit: prove knowledge of factors
        let circomCode = """
        pragma circom 2.1.0;

        template Multiplier() {
            signal input a;
            signal input b;
            signal output c;

            c <== a * b;
        }

        component main = Multiplier();
        """

        // Compile the circuit
        print("Compiling circuit...")
        let circuit = try await zk.circuits.compile(
            CircuitSource(
                code: circomCode,
                language: .circom
            )
        )

        print("Circuit compiled:")
        print("  Circuit ID: \(circuit.circuitId)")
        print("  Constraints: \(circuit.constraints)")
        print("  Public inputs: \(circuit.publicInputs)")
        print("  Private inputs: \(circuit.privateInputs)")
        print("  Outputs: \(circuit.outputs)")

        // Get circuit info
        let info = try await zk.circuits.get(circuitId: circuit.circuitId)
        print("\nCircuit info:")
        print("  Name: \(info.name)")
        print("  Version: \(info.version)")
        print("  Wires: \(info.wireCount)")
        print("  Labels: \(info.labelCount)")

        // List available circuits
        let circuits = try await zk.circuits.list()
        print("\nAvailable circuits: \(circuits.count)")
        for c in circuits.prefix(3) {
            print("  \(c.circuitId): \(c.name) (\(c.constraints) constraints)")
        }

        print()
    }

    static func groth16Example(zk: SynorZk) async throws {
        print("=== Groth16 Proving System ===")

        // Use a pre-compiled circuit
        let circuitId = "multiplier-v1"

        // Generate proving/verification keys (trusted setup)
        print("Generating keys...")
        let keys = try await zk.groth16.setup(circuitId: circuitId)
        print("Keys generated:")
        print("  Proving key size: \(keys.provingKey.count) bytes")
        print("  Verification key size: \(keys.verificationKey.count) bytes")

        // Prepare witness (private inputs)
        let witness: [String: String] = [
            "a": "3",
            "b": "7"
        ]

        // Generate proof
        print("\nGenerating proof...")
        let proof = try await zk.groth16.prove(
            ProveRequest(
                circuitId: circuitId,
                witness: witness,
                provingKey: keys.provingKey
            )
        )

        print("Proof generated:")
        print("  Proof size: \(proof.proofBytes.count) bytes")
        print("  Public signals: \(proof.publicSignals)")
        print("  Proving time: \(proof.provingTimeMs)ms")

        // Verify proof
        print("\nVerifying proof...")
        let verified = try await zk.groth16.verify(
            VerifyRequest(
                proof: proof.proofBytes,
                publicSignals: proof.publicSignals,
                verificationKey: keys.verificationKey
            )
        )

        print("Verification result: \(verified)")

        // Export proof for on-chain verification
        let solidityCalldata = try await zk.groth16.exportCalldata(proof: proof)
        print("\nSolidity calldata: \(solidityCalldata.prefix(100))...")

        print()
    }

    static func plonkExample(zk: SynorZk) async throws {
        print("=== PLONK Proving System ===")

        let circuitId = "multiplier-v1"

        // PLONK uses universal trusted setup
        print("Getting universal setup...")
        let srs = try await zk.plonk.getUniversalSetup(powersOfTau: 14) // 2^14 constraints
        print("SRS loaded: \(srs.count) bytes")

        // Generate circuit-specific keys
        print("\nGenerating circuit keys...")
        let keys = try await zk.plonk.setup(circuitId: circuitId, srs: srs)
        print("Keys generated")

        // Generate proof
        let witness: [String: String] = ["a": "5", "b": "9"]
        print("\nGenerating PLONK proof...")
        let proof = try await zk.plonk.prove(
            PlonkProveRequest(
                circuitId: circuitId,
                witness: witness,
                provingKey: keys.provingKey
            )
        )

        print("Proof generated:")
        print("  Proof size: \(proof.proofBytes.count) bytes")
        print("  Proving time: \(proof.provingTimeMs)ms")

        // Verify proof
        let verified = try await zk.plonk.verify(
            PlonkVerifyRequest(
                proof: proof.proofBytes,
                publicSignals: proof.publicSignals,
                verificationKey: keys.verificationKey
            )
        )
        print("Verification result: \(verified)")

        // Compare with Groth16
        print("\nPLONK advantages:")
        print("  - Universal trusted setup")
        print("  - Larger proofs (~2.5 KB vs ~200 bytes)")
        print("  - Faster proving for some circuits")

        print()
    }

    static func starkExample(zk: SynorZk) async throws {
        print("=== STARK Proving System ===")

        let circuitId = "multiplier-v1"

        // STARKs don't need trusted setup
        print("Configuring STARK parameters...")
        let config = StarkConfig(
            fieldSize: 256,
            hashFunction: .poseidon,
            blowupFactor: 8,
            numQueries: 30,
            foldingFactor: 8
        )

        // Generate proof
        let witness: [String: String] = ["a": "11", "b": "13"]
        print("\nGenerating STARK proof...")
        let proof = try await zk.stark.prove(
            StarkProveRequest(
                circuitId: circuitId,
                witness: witness,
                config: config
            )
        )

        print("Proof generated:")
        print("  Proof size: \(proof.proofBytes.count) bytes")
        print("  Proving time: \(proof.provingTimeMs)ms")
        print("  FRI layers: \(proof.friLayers)")

        // Verify proof
        print("\nVerifying STARK proof...")
        let verified = try await zk.stark.verify(
            StarkVerifyRequest(
                proof: proof.proofBytes,
                publicSignals: proof.publicSignals,
                config: config
            )
        )

        print("Verification result: \(verified)")

        // Compare with SNARKs
        print("\nSTARK advantages:")
        print("  - No trusted setup needed")
        print("  - Post-quantum secure")
        print("  - Larger proofs (~100 KB)")
        print("  - Faster proving for complex computations")

        print()
    }

    static func recursiveExample(zk: SynorZk) async throws {
        print("=== Recursive Proof Composition ===")

        // Create inner proofs
        print("Generating inner proofs...")
        var innerProofs: [RecursiveProof] = []

        for i in 1...3 {
            let witness: [String: String] = ["a": "\(i)", "b": "\(i + 1)"]
            let proof = try await zk.recursive.proveInner(
                RecursiveProveRequest(
                    circuitId: "multiplier-v1",
                    witness: witness,
                    level: 0
                )
            )
            innerProofs.append(proof)
            print("  Inner proof \(i) generated")
        }

        // Aggregate proofs recursively
        print("\nAggregating proofs...")
        let aggregatedProof = try await zk.recursive.aggregate(
            AggregateRequest(
                proofs: innerProofs,
                aggregationCircuit: "recursive-aggregator-v1"
            )
        )

        print("Aggregated proof:")
        print("  Proof size: \(aggregatedProof.proofBytes.count) bytes")
        print("  Proofs aggregated: \(aggregatedProof.proofsAggregated)")
        print("  Recursion depth: \(aggregatedProof.recursionDepth)")

        // Verify aggregated proof (verifies all inner proofs at once)
        print("\nVerifying aggregated proof...")
        let verified = try await zk.recursive.verifyAggregated(proof: aggregatedProof)
        print("Verification result: \(verified)")

        // Use cases
        print("\nRecursive proof use cases:")
        print("  - Rollup batch verification")
        print("  - Incremental computation proofs")
        print("  - Cross-chain state proofs")

        print()
    }

    static func ceremonyExample(zk: SynorZk) async throws {
        print("=== Trusted Setup Ceremony ===")

        // List active ceremonies
        let ceremonies = try await zk.ceremony.list(status: .active)
        print("Active ceremonies: \(ceremonies.count)")

        for ceremony in ceremonies.prefix(3) {
            print("\n  \(ceremony.ceremonyId):")
            print("    Circuit: \(ceremony.circuitId)")
            print("    Participants: \(ceremony.participantCount)")
            print("    Current round: \(ceremony.currentRound)")
            print("    Status: \(ceremony.status)")
        }

        // Create a new ceremony
        print("\nCreating new ceremony...")
        let newCeremony = try await zk.ceremony.create(
            CeremonyConfig(
                circuitId: "new-circuit-v1",
                minParticipants: 10,
                maxParticipants: 100,
                roundDuration: 3600, // 1 hour per round
                verifyContributions: true
            )
        )

        print("Ceremony created:")
        print("  Ceremony ID: \(newCeremony.ceremonyId)")
        print("  Join URL: \(newCeremony.joinUrl)")

        // Participate in a ceremony
        print("\nParticipating in ceremony...")
        let contribution = try await zk.ceremony.contribute(
            ContributionRequest(
                ceremonyId: newCeremony.ceremonyId,
                entropy: generateEntropy()
            )
        )

        print("Contribution made:")
        print("  Contribution ID: \(contribution.contributionId)")
        print("  Position: \(contribution.position)")
        print("  Hash: \(contribution.hash)")

        // Verify a contribution
        print("\nVerifying contribution...")
        let valid = try await zk.ceremony.verifyContribution(contributionId: contribution.contributionId)
        print("Contribution valid: \(valid)")

        // Get ceremony transcript (for auditability)
        let transcript = try await zk.ceremony.getTranscript(ceremonyId: newCeremony.ceremonyId)
        print("\nCeremony transcript:")
        print("  Total contributions: \(transcript.contributions.count)")
        print("  Start time: \(transcript.startTime)")
        print("  Final hash: \(transcript.finalHash ?? "pending")")

        print()
    }
}

// Helper function to generate random entropy
func generateEntropy() -> Data {
    var bytes = [UInt8](repeating: 0, count: 32)
    let status = SecRandomCopyBytes(kSecRandomDefault, bytes.count, &bytes)
    guard status == errSecSuccess else {
        fatalError("Failed to generate random bytes")
    }
    return Data(bytes)
}
