#!/bin/bash
set -e

echo "Starting ZKP Service..."

echo "Waiting for PostgreSQL..."
while ! nc -z postgres 5432; do
  sleep 1
done
echo "PostgreSQL is ready!"

echo "Checking Redis..."
if nc -z redis 6379; then
  echo "Redis is ready!"
else
  echo "Redis not available, continuing without cache"
fi

echo "Compile circuits if needed..."

if [ -f "circuits/license_verification.circom" ]; then
  if [ ! -f "circuits/build/license_verification_js/license_verification.wasm" ]; then
    echo "Compiling circuits..."
    cd circuits
    
    if [ ! -d "node_modules/circomlib" ]; then
      echo "Installing circomlib..."
      npm install circomlib@2.0.5
    fi
    
    echo "Running circom compiler..."
    circom license_verification.circom --r1cs --wasm --sym -o build
    
    if [ -f "build/license_verification.r1cs" ]; then
      echo "Circuits compiled successfully!"
      
      echo "Circuit information:"
      npx snarkjs r1cs info build/license_verification.r1cs
      
      ls -la build/
    else
      echo "ERROR: Circuit compilation failed!"
      exit 1
    fi
    
    cd ..
  else
    echo "Circuits already compiled"
  fi
fi

echo "Generate proving keys..."

if [ ! -f "keys/license_verification.zkey" ]; then
  echo "Setting up trusted setup (WARNING: Not secure for production!)..."
  
  # Create keys directory if it doesn't exist
  mkdir -p keys
  
  cd circuits
  
  # Get constraint count to determine required Powers of Tau size
  if [ -f "build/license_verification.r1cs" ]; then
    echo "Analyzing circuit constraints..."
    CONSTRAINTS=$(npx snarkjs r1cs info build/license_verification.r1cs 2>/dev/null | grep "# of Constraints:" | awk '{print $4}')
    echo "Circuit has $CONSTRAINTS constraints"
    
    REQUIRED=$((CONSTRAINTS * 2))
    POWER=15
    SIZE=32768
    
    while [ $SIZE -lt $REQUIRED ]; do
      POWER=$((POWER + 1))
      SIZE=$((SIZE * 2))
    done
    
    echo "Using Powers of Tau with 2^$POWER = $SIZE (required: $REQUIRED)"
  else
    echo "WARNING: Could not determine constraint count, using default size 15"
    POWER=15
  fi
  
  # Powers of tau ceremony
  PTAU_FILE="pot${POWER}_final.ptau"
  
  if [ ! -f "$PTAU_FILE" ]; then
    echo "=== Powers of Tau Ceremony (2^$POWER) ==="
    
    # if [ "$USE_PREGENERATED_PTAU" = "true" ]; then
    #   echo "Downloading pre-generated powers of tau (2^$POWER)..."
      
    #   # Hermez provides files up to 2^28, we'll use the appropriate size
    #   if [ $POWER -le 28 ]; then
    #     curl -L "https://hermez.s3-eu-west-1.amazonaws.com/powersOfTau28_hez_final_${POWER}.ptau" -o $PTAU_FILE
        
    #     if [ ! -f "$PTAU_FILE" ]; then
    #       echo "ERROR: Failed to download powers of tau for 2^$POWER"
    #       echo "Falling back to local generation..."
    #       USE_PREGENERATED_PTAU="false"
    #     else
    #       echo "Downloaded powers of tau successfully: $(ls -lh $PTAU_FILE | awk '{print $5}')"
    #     fi
    #   else
    #     echo "Powers of Tau 2^$POWER not available for download, generating locally..."
    #     USE_PREGENERATED_PTAU="false"
    #   fi
    # fi
    
    if [ "$USE_PREGENERATED_PTAU" != "true" ] || [ ! -f "$PTAU_FILE" ]; then
      echo "Generating powers of tau locally (2^$POWER)..."
      echo "NOTE: This may take several minutes for large circuits..."
      
      echo "Step 1: Creating initial powers of tau (2^$POWER)..."
      npx snarkjs powersoftau new bn128 $POWER pot${POWER}_0000.ptau
      
      if [ ! -f "pot${POWER}_0000.ptau" ]; then
        echo "ERROR: Failed to create initial ptau"
        exit 1
      fi
      echo "Initial ptau created: $(ls -lh pot${POWER}_0000.ptau | awk '{print $5}')"
      
      echo "Step 2: First contribution..."
      npx snarkjs powersoftau contribute pot${POWER}_0000.ptau pot${POWER}_0001.ptau \
        --name="First Contributor" -e="random entropy $(date +%s)"
      
      if [ ! -f "pot${POWER}_0001.ptau" ]; then
        echo "ERROR: Failed to create first contribution"
        exit 1
      fi
      echo "First contribution complete: $(ls -lh pot${POWER}_0001.ptau | awk '{print $5}')"
      
      echo "Step 3: Adding beacon..."
      npx snarkjs powersoftau beacon pot${POWER}_0001.ptau pot${POWER}_beacon.ptau \
        0102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f20 10 \
        -n="Final Beacon"
      
      if [ ! -f "pot${POWER}_beacon.ptau" ]; then
        echo "ERROR: Failed to add beacon"
        exit 1
      fi
      echo "Beacon added: $(ls -lh pot${POWER}_beacon.ptau | awk '{print $5}')"
      
      echo "Step 4: Preparing phase 2..."
      npx snarkjs powersoftau prepare phase2 pot${POWER}_beacon.ptau $PTAU_FILE -v
      
      if [ ! -f "$PTAU_FILE" ]; then
        echo "ERROR: Failed to prepare phase 2"
        exit 1
      fi
      echo "Phase 2 prepared: $(ls -lh $PTAU_FILE | awk '{print $5}')"
      
      echo "Step 5: Verifying powers of tau..."
      npx snarkjs powersoftau verify $PTAU_FILE
      
      # Cleanup intermediate files
      echo "Cleaning up intermediate files..."
      rm -f pot${POWER}_0000.ptau pot${POWER}_0001.ptau pot${POWER}_beacon.ptau
      
      echo "Powers of Tau ceremony complete!"
    fi
  else
    echo "Powers of Tau file already exists: $(ls -lh $PTAU_FILE | awk '{print $5}')"
  fi
  
  # Generate zkey
  if [ -f "build/license_verification.r1cs" ]; then
    echo ""
    echo "=== Generating Proving Keys ==="
    
    echo "Step 1: Groth16 setup..."
    echo "Using Powers of Tau: $PTAU_FILE"
    npx snarkjs groth16 setup build/license_verification.r1cs $PTAU_FILE ../keys/license_verification_0000.zkey
    
    if [ ! -f "../keys/license_verification_0000.zkey" ]; then
      echo "ERROR: Failed to generate initial zkey"
      echo "Check if Powers of Tau file is large enough for circuit constraints"
      exit 1
    fi
    echo "Initial zkey created: $(ls -lh ../keys/license_verification_0000.zkey | awk '{print $5}')"
    
    echo "Step 2: Contributing to phase 2..."
    npx snarkjs zkey contribute ../keys/license_verification_0000.zkey ../keys/license_verification_0001.zkey \
      --name="License Verification Contributor" -v -e="random entropy $(date +%s)"
    
    if [ ! -f "../keys/license_verification_0001.zkey" ]; then
      echo "ERROR: Failed to contribute to zkey"
      exit 1
    fi
    echo "Contribution complete: $(ls -lh ../keys/license_verification_0001.zkey | awk '{print $5}')"
    
    echo "Step 3: Adding final beacon..."
    npx snarkjs zkey beacon ../keys/license_verification_0001.zkey ../keys/license_verification_final.zkey \
      0102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f20 10 \
      -n="License Verification Final Beacon"
    
    if [ ! -f "../keys/license_verification_final.zkey" ]; then
      echo "ERROR: Failed to add final beacon"
      exit 1
    fi
    echo "Final beacon added: $(ls -lh ../keys/license_verification_final.zkey | awk '{print $5}')"
    
    echo "Step 4: Exporting verification key..."
    npx snarkjs zkey export verificationkey ../keys/license_verification_final.zkey \
      ../keys/license_verification_verification_key.json
    
    if [ ! -f "../keys/license_verification_verification_key.json" ]; then
      echo "ERROR: Failed to export verification key"
      exit 1
    fi
    echo "Verification key exported: $(ls -lh ../keys/license_verification_verification_key.json | awk '{print $5}')"
    
    echo "Step 5: Verifying final zkey..."
    npx snarkjs zkey verify build/license_verification.r1cs $PTAU_FILE ../keys/license_verification_final.zkey
    
    # Rename final key
    echo "Step 6: Finalizing keys..."
    mv ../keys/license_verification_final.zkey ../keys/license_verification.zkey
    rm -f ../keys/license_verification_0000.zkey ../keys/license_verification_0001.zkey
    
    echo ""
    echo "=== Trusted Setup Complete! ==="
    echo "Generated files:"
    echo "  - Proving key: keys/license_verification.zkey"
    echo "  - Verification key: keys/license_verification_verification_key.json"
    ls -lh ../keys/
  else
    echo "ERROR: No r1cs file found at circuits/build/license_verification.r1cs"
    echo "Available files in circuits/build/:"
    ls -la build/ 2>/dev/null || echo "Build directory not found"
    exit 1
  fi
  
  cd ..
else
  echo "Proving keys already exist:"
  ls -lh keys/
fi

echo ""
echo "=== ZKP Service Setup Complete ==="
echo "Starting main application..."

# Execute the main command
exec "$@"