#!/bin/bash

# License Verification ZKP - Circuit Setup Script
# This script compiles Circom circuits and generates proving/verification keys

set -e

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Configuration
CIRCUIT_NAME="license_verification"
CIRCUIT_DIR="${CIRCUIT_DIR:-/app/circuits}"
KEYS_DIR="${KEYS_DIR:-/app/keys}"
PTAU_DIR="${PTAU_DIR:-/app/ptau}"
BUILD_DIR="${CIRCUIT_DIR}/build"
CONSTRAINT_SIZE=14  # 2^14 constraints (~16k)

# Function to print colored output
print_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

print_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

print_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

print_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# Function to check if a command exists
command_exists() {
    command -v "$1" >/dev/null 2>&1
}

# Function to create directories
setup_directories() {
    print_info "Setting up directories..."
    mkdir -p "$BUILD_DIR"
    mkdir -p "$KEYS_DIR"
    mkdir -p "$PTAU_DIR"
    print_success "Directories created"
}

# Function to check prerequisites
check_prerequisites() {
    print_info "Checking prerequisites..."
    
    if ! command_exists circom; then
        print_error "Circom is not installed"
        print_info "Installing circom..."
        
        # Try to install circom
        if command_exists cargo; then
            git clone https://github.com/iden3/circom.git /tmp/circom
            cd /tmp/circom
            cargo build --release
            sudo mv target/release/circom /usr/local/bin/
            cd -
            rm -rf /tmp/circom
        else
            print_error "Rust/Cargo not found. Please install circom manually"
            exit 1
        fi
    fi
    
    if ! command_exists snarkjs; then
        print_warning "snarkjs not found globally, will use npx"
        SNARKJS="npx snarkjs"
    else
        SNARKJS="snarkjs"
    fi
    
    if ! command_exists node; then
        print_error "Node.js is not installed"
        exit 1
    fi
    
    print_success "Prerequisites check complete"
}

# Function to create the circuit file if it doesn't exist
create_circuit() {
    local circuit_file="${CIRCUIT_DIR}/${CIRCUIT_NAME}.circom"
    
    if [ ! -f "$circuit_file" ]; then
        print_info "Creating circuit file..."
        cat > "$circuit_file" << 'CIRCOM'
pragma circom 2.1.0;

include "../node_modules/circomlib/circuits/poseidon.circom";
include "../node_modules/circomlib/circuits/comparators.circom";
include "../node_modules/circomlib/circuits/mux1.circom";

template LicenseVerification() {
    // Private inputs - License data
    signal input licenseNumber;
    signal input practitionerName;
    signal input issuedDate;
    signal input expiryDate;
    signal input jurisdiction;
    
    // Private inputs - Merkle proof (depth 20)
    signal input pathElements[20];
    signal input pathIndices[20];
    
    // Public inputs
    signal input merkleRoot;
    signal input currentTimestamp;
    signal input minExpiryTimestamp;
    
    // Hash the license data using Poseidon
    component hasher = Poseidon(5);
    hasher.inputs[0] <== licenseNumber;
    hasher.inputs[1] <== practitionerName;
    hasher.inputs[2] <== issuedDate;
    hasher.inputs[3] <== expiryDate;
    hasher.inputs[4] <== jurisdiction;
    
    signal leafHash <== hasher.out;
    
    // Verify Merkle proof
    component merkleProof[20];
    component mux[20];
    
    signal computedHash[21];
    computedHash[0] <== leafHash;
    
    for (var i = 0; i < 20; i++) {
        merkleProof[i] = Poseidon(2);
        mux[i] = Mux1();
        
        mux[i].c[0] <== computedHash[i];
        mux[i].c[1] <== pathElements[i];
        mux[i].s <== pathIndices[i];
        
        merkleProof[i].inputs[0] <== mux[i].out;
        merkleProof[i].inputs[1] <== pathElements[i] - mux[i].out + computedHash[i];
        
        computedHash[i + 1] <== merkleProof[i].out;
    }
    
    // Check merkle root matches
    signal rootMatch <== computedHash[20] - merkleRoot;
    rootMatch === 0;
    
    // Check license is not expired
    component notExpired = GreaterEqThan(64);
    notExpired.in[0] <== expiryDate;
    notExpired.in[1] <== minExpiryTimestamp;
    
    // Check license was issued before current time
    component wasIssued = LessEqThan(64);
    wasIssued.in[0] <== issuedDate;
    wasIssued.in[1] <== currentTimestamp;
    
    // Output signals
    signal output isValid;
    isValid <== notExpired.out * wasIssued.out;
}

component main = LicenseVerification();
CIRCOM
        print_success "Circuit file created"
    else
        print_info "Circuit file already exists"
    fi
}

# Function to compile the circuit
compile_circuit() {
    print_info "Compiling circuit '${CIRCUIT_NAME}'..."
    
    cd "$CIRCUIT_DIR"
    
    # Check if circomlib is installed
    if [ ! -d "node_modules/circomlib" ]; then
        print_info "Installing circomlib..."
        npm install circomlib@2.0.5
    fi
    
    # Compile with circom
    print_info "Running circom compiler..."
    circom "${CIRCUIT_NAME}.circom" \
        --r1cs \
        --wasm \
        --sym \
        --c \
        -o "$BUILD_DIR" \
        -l node_modules
    
    if [ -f "${BUILD_DIR}/${CIRCUIT_NAME}.r1cs" ]; then
        print_success "Circuit compiled successfully"
        
        # Print circuit info
        print_info "Circuit statistics:"
        $SNARKJS r1cs info "${BUILD_DIR}/${CIRCUIT_NAME}.r1cs"
    else
        print_error "Circuit compilation failed"
        exit 1
    fi
}

# Function to download or generate Powers of Tau
setup_powers_of_tau() {
    local ptau_file="${PTAU_DIR}/pot${CONSTRAINT_SIZE}_final.ptau"
    
    if [ -f "$ptau_file" ]; then
        print_info "Powers of Tau file already exists"
        return 0
    fi
    
    print_info "Setting up Powers of Tau..."
    
    # Option 1: Download from Hermez ceremony (recommended for production)
    if [ "$USE_HERMEZ_PTAU" = "true" ]; then
        print_info "Downloading Powers of Tau from Hermez ceremony..."
        curl -L "https://hermez.s3-eu-west-1.amazonaws.com/powersOfTau28_hez_final_${CONSTRAINT_SIZE}.ptau" \
             -o "$ptau_file"
        print_success "Downloaded trusted Powers of Tau"
    else
        # Option 2: Generate locally (for testing only!)
        print_warning "Generating Powers of Tau locally (NOT SECURE FOR PRODUCTION!)"
        
        # Start new ceremony
        $SNARKJS powersoftau new bn128 ${CONSTRAINT_SIZE} "${PTAU_DIR}/pot${CONSTRAINT_SIZE}_0000.ptau"
        
        # Contribute to ceremony
        $SNARKJS powersoftau contribute \
            "${PTAU_DIR}/pot${CONSTRAINT_SIZE}_0000.ptau" \
            "${PTAU_DIR}/pot${CONSTRAINT_SIZE}_0001.ptau" \
            --name="First contribution" \
            -v -e="$(head -n 1024 /dev/urandom | sha256sum | head -c 64)"
        
        # Add beacon
        $SNARKJS powersoftau beacon \
            "${PTAU_DIR}/pot${CONSTRAINT_SIZE}_0001.ptau" \
            "${PTAU_DIR}/pot${CONSTRAINT_SIZE}_beacon.ptau" \
            "0102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f20" 10 \
            -n="Final Beacon"
        
        # Prepare phase 2
        $SNARKJS powersoftau prepare phase2 \
            "${PTAU_DIR}/pot${CONSTRAINT_SIZE}_beacon.ptau" \
            "$ptau_file" \
            -v
        
        # Cleanup intermediate files
        rm -f "${PTAU_DIR}/pot${CONSTRAINT_SIZE}_0000.ptau"
        rm -f "${PTAU_DIR}/pot${CONSTRAINT_SIZE}_0001.ptau"
        rm -f "${PTAU_DIR}/pot${CONSTRAINT_SIZE}_beacon.ptau"
        
        print_success "Generated Powers of Tau (TEST ONLY)"
    fi
    
    # Verify the ptau file
    print_info "Verifying Powers of Tau..."
    $SNARKJS powersoftau verify "$ptau_file"
}

# Function to generate proving and verification keys
generate_keys() {
    print_info "Generating proving and verification keys..."
    
    local r1cs_file="${BUILD_DIR}/${CIRCUIT_NAME}.r1cs"
    local ptau_file="${PTAU_DIR}/pot${CONSTRAINT_SIZE}_final.ptau"
    local zkey_0="${KEYS_DIR}/${CIRCUIT_NAME}_0000.zkey"
    local zkey_1="${KEYS_DIR}/${CIRCUIT_NAME}_0001.zkey"
    local zkey_final="${KEYS_DIR}/${CIRCUIT_NAME}.zkey"
    local vkey_file="${KEYS_DIR}/${CIRCUIT_NAME}_verification_key.json"
    
    # Check if keys already exist
    if [ -f "$zkey_final" ] && [ -f "$vkey_file" ]; then
        print_warning "Keys already exist. Use --force to regenerate"
        if [ "$1" != "--force" ]; then
            return 0
        fi
    fi
    
    # Setup groth16
    print_info "Running Groth16 setup..."
    $SNARKJS groth16 setup \
        "$r1cs_file" \
        "$ptau_file" \
        "$zkey_0"
    
    # Contribute to the ceremony
    print_info "Contributing to phase 2 ceremony..."
    $SNARKJS zkey contribute \
        "$zkey_0" \
        "$zkey_1" \
        --name="License Verification Contribution" \
        -v -e="$(date +%s | sha256sum | head -c 64)"
    
    # Add beacon
    print_info "Adding beacon to finalize..."
    $SNARKJS zkey beacon \
        "$zkey_1" \
        "$zkey_final" \
        "0102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f20" 10 \
        -n="License Verification Final Beacon"
    
    # Export verification key
    print_info "Exporting verification key..."
    $SNARKJS zkey export verificationkey \
        "$zkey_final" \
        "$vkey_file"
    
    # Verify the final zkey
    print_info "Verifying final zkey..."
    $SNARKJS zkey verify \
        "$r1cs_file" \
        "$ptau_file" \
        "$zkey_final"
    
    # Cleanup intermediate files
    rm -f "$zkey_0" "$zkey_1"
    
    print_success "Proving and verification keys generated"
    
    # Print key sizes
    print_info "Key file sizes:"
    ls -lh "$zkey_final" "$vkey_file" | awk '{print "  " $9 ": " $5}'
}

# Function to generate a test proof
test_proof() {
    print_info "Generating test proof..."
    
    local wasm_file="${BUILD_DIR}/${CIRCUIT_NAME}_js/${CIRCUIT_NAME}.wasm"
    local zkey_file="${KEYS_DIR}/${CIRCUIT_NAME}.zkey"
    local vkey_file="${KEYS_DIR}/${CIRCUIT_NAME}_verification_key.json"
    
    # Create test input
    cat > /tmp/test_input.json << JSON
{
    "licenseNumber": "12345678901234567890",
    "practitionerName": "98765432109876543210",
    "issuedDate": "1609459200",
    "expiryDate": "1735689600",
    "jurisdiction": "11111111111111111111",
    "pathElements": [
        "1", "2", "3", "4", "5", "6", "7", "8", "9", "10",
        "11", "12", "13", "14", "15", "16", "17", "18", "19", "20"
    ],
    "pathIndices": [
        "0", "1", "0", "1", "0", "1", "0", "1", "0", "1",
        "0", "1", "0", "1", "0", "1", "0", "1", "0", "1"
    ],
    "merkleRoot": "123456789",
    "currentTimestamp": "1609459200",
    "minExpiryTimestamp": "1609545600"
}
JSON
    
    # Generate witness
    print_info "Calculating witness..."
    node "${BUILD_DIR}/${CIRCUIT_NAME}_js/generate_witness.js" \
        "$wasm_file" \
        /tmp/test_input.json \
        /tmp/witness.wtns
    
    # Generate proof
    print_info "Generating proof..."
    local start_time=$(date +%s%N)
    
    $SNARKJS groth16 prove \
        "$zkey_file" \
        /tmp/witness.wtns \
        /tmp/proof.json \
        /tmp/public.json
    
    local end_time=$(date +%s%N)
    local duration=$((($end_time - $start_time) / 1000000))
    
    print_success "Proof generated in ${duration}ms"
    
    # Verify proof
    print_info "Verifying proof..."
    $SNARKJS groth16 verify \
        "$vkey_file" \
        /tmp/public.json \
        /tmp/proof.json
    
    if [ $? -eq 0 ]; then
        print_success "Proof verified successfully!"
    else
        print_error "Proof verification failed"
        exit 1
    fi
    
    # Cleanup
    rm -f /tmp/test_input.json /tmp/witness.wtns /tmp/proof.json /tmp/public.json
}

# Function to display summary
display_summary() {
    echo ""
    echo "========================================="
    echo -e "${GREEN}Circuit Setup Complete!${NC}"
    echo "========================================="
    echo ""
    echo "Generated files:"
    echo "  Circuit: ${BUILD_DIR}/${CIRCUIT_NAME}.r1cs"
    echo "  WASM:    ${BUILD_DIR}/${CIRCUIT_NAME}_js/${CIRCUIT_NAME}.wasm"
    echo "  ZKey:    ${KEYS_DIR}/${CIRCUIT_NAME}.zkey"
    echo "  VKey:    ${KEYS_DIR}/${CIRCUIT_NAME}_verification_key.json"
    echo ""
    
    # Get file sizes
    if [ -f "${KEYS_DIR}/${CIRCUIT_NAME}.zkey" ]; then
        local zkey_size=$(ls -lh "${KEYS_DIR}/${CIRCUIT_NAME}.zkey" | awk '{print $5}')
        local vkey_size=$(ls -lh "${KEYS_DIR}/${CIRCUIT_NAME}_verification_key.json" | awk '{print $5}')
        echo "File sizes:"
        echo "  ZKey: $zkey_size"
        echo "  VKey: $vkey_size"
        echo ""
    fi
    
    echo "Next steps:"
    echo "  1. Start the ZKP service: npm start"
    echo "  2. Run benchmarks: npm run benchmark"
    echo "  3. Access test UI: http://localhost:3000"
    echo ""
}

# Main execution
main() {
    echo "========================================="
    echo "License Verification Circuit Setup"
    echo "========================================="
    echo ""
    
    # Parse arguments
    FORCE_REBUILD=false
    SKIP_TEST=false
    
    for arg in "$@"; do
        case $arg in
            --force)
                FORCE_REBUILD=true
                print_warning "Force rebuild enabled"
                ;;
            --skip-test)
                SKIP_TEST=true
                ;;
            --use-hermez)
                USE_HERMEZ_PTAU=true
                ;;
            --help)
                echo "Usage: $0 [OPTIONS]"
                echo ""
                echo "Options:"
                echo "  --force        Force rebuild of all components"
                echo "  --skip-test    Skip test proof generation"
                echo "  --use-hermez   Use Hermez trusted setup (recommended)"
                echo "  --help         Show this help message"
                exit 0
                ;;
        esac
    done
    
    # Run setup steps
    check_prerequisites
    setup_directories
    create_circuit
    compile_circuit
    setup_powers_of_tau
    generate_keys $FORCE_REBUILD
    
    # Test the setup
    if [ "$SKIP_TEST" != "true" ]; then
        test_proof
    fi
    
    display_summary
}

# Run main function
main "$@"