#!/usr/bin/env bash
# MEV Bot Keystore Setup Script
# Encrypts and stores the private key securely

set -euo pipefail

KEYSTORE_DIR="keystore/production"
PRIVATE_KEY_FILE="/tmp/wallet_key.txt"
ENCRYPTION_KEY_ENV="MEV_BOT_ENCRYPTION_KEY"

echo "═══════════════════════════════════════════════════════════"
echo "🔐 MEV Bot Keystore Configuration"
echo "═══════════════════════════════════════════════════════════"
echo ""

# Check if private key file exists
if [ ! -f "$PRIVATE_KEY_FILE" ]; then
    echo "❌ Error: Private key file not found at $PRIVATE_KEY_FILE"
    echo "   Please ensure the wallet has been generated first."
    exit 1
fi

# Check if encryption key is set
ENCRYPTION_KEY="${MEV_BOT_ENCRYPTION_KEY:-}"
if [ -z "$ENCRYPTION_KEY" ]; then
    echo "⚠️  Warning: $ENCRYPTION_KEY_ENV not set in environment"
    echo ""
    echo "📝 Setting up encryption key..."

    # Generate a secure encryption key
    ENCRYPTION_KEY=$(openssl rand -base64 32)
    echo "export $ENCRYPTION_KEY_ENV=\"$ENCRYPTION_KEY\"" >> .env.production
    export MEV_BOT_ENCRYPTION_KEY="$ENCRYPTION_KEY"

    echo "✅ Generated and saved encryption key to .env.production"
fi

# Create keystore directory
mkdir -p "$KEYSTORE_DIR"

# Read private key
PRIVATE_KEY=$(cat "$PRIVATE_KEY_FILE")

# Derive wallet address from private key (using cast if available, otherwise Python)
echo ""
echo "🔍 Deriving wallet address from private key..."

if command -v cast &> /dev/null; then
    # Use Foundry's cast tool
    WALLET_ADDRESS=$(cast wallet address "$PRIVATE_KEY")
    echo "✅ Wallet Address: $WALLET_ADDRESS"
else
    # Use Python with web3.py (if available)
    if command -v python3 &> /dev/null; then
        WALLET_ADDRESS=$(python3 << EOF
from eth_account import Account
import sys

try:
    private_key = "$PRIVATE_KEY"
    if private_key.startswith('0x'):
        private_key = private_key[2:]

    account = Account.from_key(bytes.fromhex(private_key))
    print(account.address)
except Exception as e:
    print(f"Error: {e}", file=sys.stderr)
    sys.exit(1)
EOF
)
        if [ $? -eq 0 ]; then
            echo "✅ Wallet Address: $WALLET_ADDRESS"
        else
            echo "⚠️  Warning: Could not derive address automatically"
            echo "   Please provide your wallet address manually."
            read -p "Enter wallet address (from MetaMask): " WALLET_ADDRESS
        fi
    else
        echo "⚠️  Neither cast nor python3 available"
        echo "   Please provide your wallet address manually."
        read -p "Enter wallet address (from MetaMask): " WALLET_ADDRESS
    fi
fi

# Create keystore file with encryption
KEYSTORE_FILE="$KEYSTORE_DIR/executor_wallet.json"

echo ""
echo "🔐 Creating encrypted keystore..."

# Encrypt private key with AES-256-CBC
ENCRYPTED_KEY=$(echo "$PRIVATE_KEY" | openssl enc -aes-256-cbc -a -salt -pass pass:"${MEV_BOT_ENCRYPTION_KEY}")

# Create keystore JSON
cat > "$KEYSTORE_FILE" << EOF
{
  "version": 1,
  "address": "$WALLET_ADDRESS",
  "crypto": {
    "cipher": "aes-256-cbc",
    "ciphertext": "$ENCRYPTED_KEY"
  },
  "id": "$(uuidgen 2>/dev/null || echo "mev-executor-$(date +%s)")",
  "metadata": {
    "created": "$(date -u +"%Y-%m-%dT%H:%M:%SZ")",
    "purpose": "MEV Bot Executor Wallet",
    "network": "Arbitrum One (Chain ID: 42161)"
  }
}
EOF

chmod 600 "$KEYSTORE_FILE"

echo "✅ Encrypted keystore created: $KEYSTORE_FILE"
echo ""

# Verify keystore
echo "🔍 Verifying keystore integrity..."

# Test decryption
DECRYPTED_KEY=$(echo "$ENCRYPTED_KEY" | openssl enc -aes-256-cbc -d -a -pass pass:"${MEV_BOT_ENCRYPTION_KEY}")

if [ "$DECRYPTED_KEY" = "$PRIVATE_KEY" ]; then
    echo "✅ Keystore verification successful!"
else
    echo "❌ Error: Keystore verification failed!"
    echo "   Encryption/decryption mismatch detected."
    exit 1
fi

echo ""
echo "═══════════════════════════════════════════════════════════"
echo "✅ Keystore Configuration Complete!"
echo "═══════════════════════════════════════════════════════════"
echo ""
echo "📋 Summary:"
echo "  • Wallet Address: $WALLET_ADDRESS"
echo "  • Keystore File: $KEYSTORE_FILE"
echo "  • Encryption: AES-256-CBC"
echo "  • Network: Arbitrum One"
echo ""
echo "🔒 Security:"
echo "  • Private key encrypted with MEV_BOT_ENCRYPTION_KEY"
echo "  • Keystore file permissions: 600 (owner read/write only)"
echo "  • Original key file: $PRIVATE_KEY_FILE (keep secure!)"
echo ""
echo "⏭️  Next Steps:"
echo "  1. Verify wallet is funded (use check-wallet-balance.sh)"
echo "  2. Enable execution mode in bot config"
echo "  3. Restart bot to begin live trading"
echo ""