#!/usr/bin/env bash
# Configure Besu nodes for ChainID 138 with static-nodes.json and permissioned-nodes.json
# Applies to all Besu nodes including new containers: 1504 (besu-sentry-5), 2503 (besu-rpc-4)
#
# This script:
# 1. Collects enodes from all Besu nodes
# 2. Generates static-nodes.json and permissioned-nodes.json
# 3. Deploys configurations to all Besu containers
# 4. Configures discovery settings (disabled for RPC nodes that report chainID 0x1 to MetaMask for wallet compatibility)
# 5. Restarts Besu services

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)"

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

log_info() { echo -e "${BLUE}[INFO]${NC} $1"; }
log_success() { echo -e "${GREEN}[✓]${NC} $1"; }
log_warn() { echo -e "${YELLOW}[WARN]${NC} $1"; }
log_error() { echo -e "${RED}[ERROR]${NC} $1"; }

# Configuration
PROXMOX_HOST="${PROXMOX_HOST:-192.168.11.10}"
WORK_DIR="${WORK_DIR:-/tmp/besu-chain138-config}"
OUTPUT_DIR="${OUTPUT_DIR:-$PROJECT_ROOT/output/chain138-config}"

# All Besu nodes for ChainID 138
# Validators: 1000-1004
# Sentries: 1500-1503, 1504 (new)
# RPC: 2500-2502, 2503 (new)
declare -A BESU_NODES=(
    # Validators
    [1000]="192.168.11.100"
    [1001]="192.168.11.101"
    [1002]="192.168.11.102"
    [1003]="192.168.11.103"
    [1004]="192.168.11.104"
    # Sentries
    [1500]="192.168.11.150"
    [1501]="192.168.11.151"
    [1502]="192.168.11.152"
    [1503]="192.168.11.153"
    [1504]="192.168.11.154"  # New: besu-sentry-5
    # RPC Nodes
    [2500]="192.168.11.250"
    [2501]="192.168.11.251"
    [2502]="192.168.11.252"
    [2503]="192.168.11.253"  # Ali's RPC node (0x8a identity)
    [2504]="192.168.11.254"  # Ali's RPC node (0x1 identity)
    [2505]="192.168.11.255"  # Luis's RPC node (0x8a identity)
    [2506]="192.168.11.256"  # Luis's RPC node (0x1 identity)
    [2507]="192.168.11.257"  # Putu's RPC node (0x8a identity)
    [2508]="192.168.11.258"  # Putu's RPC node (0x1 identity)
)

# RPC nodes that should have discovery disabled (report chainID 0x1 to MetaMask for wallet compatibility)
# These nodes are connected to ChainID 138 but report 0x1 to MetaMask wallets
# Discovery is disabled to prevent actual connection to Ethereum mainnet while reporting 0x1 to wallets
DISCOVERY_DISABLED_VMIDS=(2500 2503 2504 2505 2506 2507 2508)

# Besu data paths
BESU_DATA_PATH="${BESU_DATA_PATH:-/var/lib/besu}"
BESU_PERMISSIONS_PATH="${BESU_PERMISSIONS_PATH:-/var/lib/besu/permissions}"
BESU_GENESIS_PATH="${BESU_GENESIS_PATH:-/genesis}"

# Create work directories
mkdir -p "$WORK_DIR"
mkdir -p "$OUTPUT_DIR"

# Function to check if container exists and is running
check_container() {
    local vmid=$1
    if ! ssh -o StrictHostKeyChecking=accept-new "root@${PROXMOX_HOST}" "pct status $vmid 2>/dev/null | grep -q running"; then
        return 1
    fi
    return 0
}

# Function to extract enode from container
extract_enode() {
    local vmid=$1
    local ip=$2
    
    log_info "Extracting enode from VMID $vmid - IP: $ip"
    
    # Method 1: Try RPC admin_nodeInfo (if RPC is enabled with ADMIN API)
    local enode_rpc
    enode_rpc=$(ssh -o StrictHostKeyChecking=accept-new "root@${PROXMOX_HOST}" "pct exec $vmid -- curl -s -X POST -H 'Content-Type: application/json' -d '{\"jsonrpc\":\"2.0\",\"method\":\"admin_nodeInfo\",\"params\":[],\"id\":1}' http://localhost:8545 2>/dev/null | python3 -c \"import sys, json; data=json.load(sys.stdin); print(data.get('result', {}).get('enode', ''))\" 2>/dev/null || echo \"\"")
    
    if [[ -n "$enode_rpc" ]] && [[ "$enode_rpc" != "null" ]] && [[ "$enode_rpc" != "" ]]; then
        # Replace IP in enode with actual IP
        echo "$enode_rpc" | sed "s/@[^:]*:/@${ip}:/"
        return 0
    fi
    
    # Method 2: Extract from nodekey using Besu CLI
    local nodekey_path
    nodekey_path=$(ssh -o StrictHostKeyChecking=accept-new "root@${PROXMOX_HOST}" "pct exec $vmid -- find ${BESU_DATA_PATH} -name 'nodekey' -o -name 'key' 2>/dev/null | head -1" || echo "")
    
    if [[ -z "$nodekey_path" ]]; then
        # Try common locations
        for path in "${BESU_DATA_PATH}/key" "${BESU_DATA_PATH}/nodekey" "/data/besu/nodekey" "/data/besu/key"; do
            if ssh -o StrictHostKeyChecking=accept-new "root@${PROXMOX_HOST}" "pct exec $vmid -- test -f $path 2>/dev/null"; then
                nodekey_path="$path"
                break
            fi
        done
    fi
    
    if [[ -n "$nodekey_path" ]]; then
        local enode
        enode=$(ssh -o StrictHostKeyChecking=accept-new "root@${PROXMOX_HOST}" "pct exec $vmid -- /opt/besu/bin/besu public-key export --node-private-key-file=\"$nodekey_path\" --format=enode 2>/dev/null" || echo "")
        
        if [[ -n "$enode" ]]; then
            # Replace IP in enode with actual IP
            echo "$enode" | sed "s/@[0-9.]*:/@${ip}:/"
            return 0
        fi
    fi
    
    log_warn "Could not extract enode for VMID $vmid"
    return 1
}

# Function to collect all enodes
collect_enodes() {
    log_info "=== Collecting Enodes from All Besu Nodes ==="
    
    local collected_file="$WORK_DIR/collected-enodes.txt"
    > "$collected_file"
    
    local success_count=0
    local fail_count=0
    
    for vmid in "${!BESU_NODES[@]}"; do
        local ip="${BESU_NODES[$vmid]}"
        
        if ! check_container "$vmid"; then
            log_warn "Container $vmid not running, skipping..."
            ((fail_count++))
            continue
        fi
        
        if enode=$(extract_enode "$vmid" "$ip"); then
            echo "$enode" >> "$collected_file"
            log_success "VMID $vmid: $enode"
            ((success_count++))
        else
            log_warn "Failed to extract enode from VMID $vmid"
            ((fail_count++))
        fi
    done
    
    log_info "Collected $success_count enodes, $fail_count failed"
    
    if [[ $success_count -eq 0 ]]; then
        log_error "No enodes collected. Cannot proceed."
        exit 1
    fi
    
    # Sort and deduplicate
    sort -u "$collected_file" > "$collected_file.tmp"
    mv "$collected_file.tmp" "$collected_file"
    
    echo "$collected_file"
}

# Function to generate static-nodes.json
generate_static_nodes() {
    local enodes_file=$1
    local output_file="$OUTPUT_DIR/static-nodes.json"
    
    log_info "=== Generating static-nodes.json ==="
    
    # Create JSON array from enodes
    python3 << PYEOF
import json
import sys

if len(sys.argv) < 3:
    print("Usage: python3 script.py <enodes_file> <output_file>", file=sys.stderr)
    sys.exit(1)

enodes_file = sys.argv[1]
output_file = sys.argv[2]

enodes = []
try:
    with open(enodes_file, 'r') as f:
        for line in f:
            enode = line.strip()
            if enode and enode.startswith('enode://'):
                enodes.append(enode)
except FileNotFoundError:
    print(f"Error: File not found: {enodes_file}", file=sys.stderr)
    sys.exit(1)

# Sort for consistency
enodes.sort()

with open(output_file, 'w') as f:
    json.dump(enodes, f, indent=2)

print("Generated static-nodes.json with {} nodes".format(len(enodes)))
PYEOF
 "$enodes_file" "$output_file"
    
    log_success "Generated: $output_file"
    echo "$output_file"
}

# Function to generate permissioned-nodes.json
generate_permissioned_nodes() {
    local enodes_file=$1
    local output_file="$OUTPUT_DIR/permissioned-nodes.json"
    
    log_info "=== Generating permissioned-nodes.json ==="
    
    # Same content as static-nodes.json (all nodes must be permissioned)
    cp "$OUTPUT_DIR/static-nodes.json" "$output_file"
    
    log_success "Generated: $output_file"
    echo "$output_file"
}

# Function to deploy files to container
deploy_to_container() {
    local vmid=$1
    local static_nodes_file=$2
    local permissioned_nodes_file=$3
    
    log_info "Deploying to VMID $vmid..."
    
    # Create directories if they don't exist
    ssh -o StrictHostKeyChecking=accept-new "root@${PROXMOX_HOST}" "pct exec $vmid -- mkdir -p ${BESU_DATA_PATH} ${BESU_PERMISSIONS_PATH} ${BESU_GENESIS_PATH} 2>/dev/null || true"
    
    # Copy static-nodes.json
    scp -o StrictHostKeyChecking=accept-new \
        "$static_nodes_file" \
        "root@${PROXMOX_HOST}:/tmp/static-nodes.json"
    
    ssh -o StrictHostKeyChecking=accept-new "root@${PROXMOX_HOST}" "pct push $vmid /tmp/static-nodes.json ${BESU_DATA_PATH}/static-nodes.json && pct push $vmid /tmp/static-nodes.json ${BESU_GENESIS_PATH}/static-nodes.json && pct exec $vmid -- chown -R besu:besu ${BESU_DATA_PATH} ${BESU_GENESIS_PATH} 2>/dev/null || pct exec $vmid -- chown -R root:root ${BESU_DATA_PATH} ${BESU_GENESIS_PATH} 2>/dev/null || true"
    
    # Copy permissioned-nodes.json
    scp -o StrictHostKeyChecking=accept-new \
        "$permissioned_nodes_file" \
        "root@${PROXMOX_HOST}:/tmp/permissioned-nodes.json"
    
    ssh -o StrictHostKeyChecking=accept-new "root@${PROXMOX_HOST}" "pct push $vmid /tmp/permissioned-nodes.json ${BESU_PERMISSIONS_PATH}/permissioned-nodes.json && pct exec $vmid -- chown -R besu:besu ${BESU_PERMISSIONS_PATH} 2>/dev/null || pct exec $vmid -- chown -R root:root ${BESU_PERMISSIONS_PATH} 2>/dev/null || true"
    
    # Set permissions
    ssh -o StrictHostKeyChecking=accept-new "root@${PROXMOX_HOST}" "pct exec $vmid -- chmod 644 ${BESU_DATA_PATH}/static-nodes.json ${BESU_GENESIS_PATH}/static-nodes.json ${BESU_PERMISSIONS_PATH}/permissioned-nodes.json 2>/dev/null || true"
    
    log_success "Deployed to VMID $vmid"
}

# Function to configure discovery settings
configure_discovery() {
    local vmid=$1
    local disable=$2
    
    log_info "Configuring discovery for VMID $vmid - disable=${disable}"
    
    # Find Besu config file
    local config_files=(
        "/etc/besu/config.toml"
        "/etc/besu/config-rpc-core.toml"
        "/etc/besu/config-rpc-perm.toml"
        "/etc/besu/config-rpc-public.toml"
        "/etc/besu/config-validator.toml"
        "/etc/besu/config-sentry.toml"
    )
    
    local config_found=false
    for config_file in "${config_files[@]}"; do
        if ssh -o StrictHostKeyChecking=accept-new "root@${PROXMOX_HOST}" "pct exec $vmid -- test -f $config_file 2>/dev/null"; then
            config_found=true
            
            if [[ "$disable" == "true" ]]; then
                # Disable discovery
                ssh -o StrictHostKeyChecking=accept-new "root@${PROXMOX_HOST}" "pct exec $vmid -- sed -i 's/^discovery-enabled=.*/discovery-enabled=false/' $config_file && sed -i 's/^# discovery-enabled=.*/discovery-enabled=false/' $config_file || echo 'discovery-enabled=false' >> $config_file"
                log_success "Discovery disabled in $config_file"
            else
                # Enable discovery (default)
                ssh -o StrictHostKeyChecking=accept-new "root@${PROXMOX_HOST}" "pct exec $vmid -- sed -i 's/^discovery-enabled=.*/discovery-enabled=true/' $config_file"
                log_success "Discovery enabled in $config_file"
            fi
            break
        fi
    done
    
    # Also check systemd service file for CLI flags
    if ssh -o StrictHostKeyChecking=accept-new "root@${PROXMOX_HOST}" \
        "pct exec $vmid -- test -f /etc/systemd/system/besu*.service 2>/dev/null"; then
        local service_file
        service_file=$(ssh -o StrictHostKeyChecking=accept-new "root@${PROXMOX_HOST}" "pct exec $vmid -- find /etc/systemd/system -name 'besu*.service' | head -1")
        
        if [[ "$disable" == "true" ]]; then
            # Add or update discovery flag
            ssh -o StrictHostKeyChecking=accept-new "root@${PROXMOX_HOST}" "pct exec $vmid -- sed -i 's/--discovery-enabled=[^ ]*/--discovery-enabled=false/' $service_file && grep -q '--discovery-enabled' $service_file || sed -i 's|ExecStart=|ExecStart=\\n  --discovery-enabled=false |' $service_file"
        else
            # Remove or update discovery flag
            ssh -o StrictHostKeyChecking=accept-new "root@${PROXMOX_HOST}" "pct exec $vmid -- sed -i 's/--discovery-enabled=false/--discovery-enabled=true/' $service_file || sed -i '/--discovery-enabled=false/d' $service_file"
        fi
        
        ssh -o StrictHostKeyChecking=accept-new "root@${PROXMOX_HOST}" "pct exec $vmid -- systemctl daemon-reload 2>/dev/null || true"
    fi
    
    if [[ "$config_found" == "false" ]]; then
        log_warn "No config file found for VMID $vmid, discovery setting may need manual configuration"
    fi
}

# Function to update Besu config to use correct file paths
update_config_paths() {
    local vmid=$1
    
    log_info "Updating config paths for VMID $vmid..."
    
    # Find Besu config file
    local config_files=(
        "/etc/besu/config.toml"
        "/etc/besu/config-rpc-core.toml"
        "/etc/besu/config-rpc-perm.toml"
        "/etc/besu/config-rpc-public.toml"
        "/etc/besu/config-validator.toml"
        "/etc/besu/config-sentry.toml"
    )
    
    for config_file in "${config_files[@]}"; do
        if ssh -o StrictHostKeyChecking=accept-new "root@${PROXMOX_HOST}" "pct exec $vmid -- test -f $config_file 2>/dev/null"; then
            # Update static-nodes-file path
            ssh -o StrictHostKeyChecking=accept-new "root@${PROXMOX_HOST}" "pct exec $vmid -- sed -i 's|^static-nodes-file=.*|static-nodes-file=\"${BESU_DATA_PATH}/static-nodes.json\"|' $config_file || sed -i 's|^static-nodes-file=.*|static-nodes-file=\"${BESU_GENESIS_PATH}/static-nodes.json\"|' $config_file || echo 'static-nodes-file=\"${BESU_DATA_PATH}/static-nodes.json\"' >> $config_file"
            
            # Update permissions-nodes-config-file path
            ssh -o StrictHostKeyChecking=accept-new "root@${PROXMOX_HOST}" "pct exec $vmid -- sed -i 's|^permissions-nodes-config-file=.*|permissions-nodes-config-file=\"${BESU_PERMISSIONS_PATH}/permissioned-nodes.json\"|' $config_file || echo 'permissions-nodes-config-file=\"${BESU_PERMISSIONS_PATH}/permissioned-nodes.json\"' >> $config_file"
            
            # Ensure permissions are enabled
            ssh -o StrictHostKeyChecking=accept-new "root@${PROXMOX_HOST}" "pct exec $vmid -- sed -i 's/^permissions-nodes-config-file-enabled=.*/permissions-nodes-config-file-enabled=true/' $config_file || echo 'permissions-nodes-config-file-enabled=true' >> $config_file"
            
            log_success "Updated paths in $config_file"
            break
        fi
    done
}

# Function to restart Besu service
restart_besu_service() {
    local vmid=$1
    
    log_info "Restarting Besu service on VMID $vmid..."
    
    # Find service name
    local service_name
    service_name=$(ssh -o StrictHostKeyChecking=accept-new "root@${PROXMOX_HOST}" "pct exec $vmid -- systemctl list-units --type=service | grep -i besu | awk '{print \$1}' | head -1" || echo "")
    
    if [[ -z "$service_name" ]]; then
        # Try common service names
        for name in "besu" "besu-rpc" "besu-validator" "besu-sentry"; do
            if ssh -o StrictHostKeyChecking=accept-new "root@${PROXMOX_HOST}" "pct exec $vmid -- systemctl list-units --type=service | grep -q $name"; then
                service_name="$name.service"
                break
            fi
        done
    fi
    
    if [[ -n "$service_name" ]]; then
        ssh -o StrictHostKeyChecking=accept-new "root@${PROXMOX_HOST}" "pct exec $vmid -- systemctl restart $service_name 2>/dev/null || true"
        sleep 3
        
        if ssh -o StrictHostKeyChecking=accept-new "root@${PROXMOX_HOST}" "pct exec $vmid -- systemctl is-active --quiet $service_name 2>/dev/null"; then
            log_success "Service $service_name restarted successfully"
        else
            log_warn "Service $service_name may not be running"
        fi
    else
        log_warn "Could not find Besu service for VMID $vmid"
    fi
}

# Main execution
main() {
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo "Besu ChainID 138 Configuration Script"
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo ""
    
    # Step 1: Collect enodes
    local enodes_file
    enodes_file=$(collect_enodes)
    
    # Step 2: Generate configuration files
    local static_nodes_file
    static_nodes_file=$(generate_static_nodes "$enodes_file")
    
    local permissioned_nodes_file
    permissioned_nodes_file=$(generate_permissioned_nodes "$enodes_file")
    
    # Step 3: Deploy to all containers
    log_info "=== Deploying Configurations to All Besu Nodes ==="
    
    for vmid in "${!BESU_NODES[@]}"; do
        if ! check_container "$vmid"; then
            log_warn "Container $vmid not running, skipping..."
            continue
        fi
        
        echo ""
        echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
        log_info "Processing VMID $vmid"
        echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
        
        # Deploy files
        deploy_to_container "$vmid" "$static_nodes_file" "$permissioned_nodes_file"
        
        # Update config paths
        update_config_paths "$vmid"
        
        # Configure discovery (disable for specific RPC nodes)
        local disable_discovery=false
        for disabled_vmid in "${DISCOVERY_DISABLED_VMIDS[@]}"; do
            if [[ "$vmid" == "$disabled_vmid" ]]; then
                disable_discovery=true
                break
            fi
        done
        
        configure_discovery "$vmid" "$disable_discovery"
        
        # Restart service
        restart_besu_service "$vmid"
    done
    
    # Cleanup
    ssh -o StrictHostKeyChecking=accept-new "root@${PROXMOX_HOST}" "rm -f /tmp/static-nodes.json /tmp/permissioned-nodes.json" 2>/dev/null || true
    
    echo ""
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    log_success "Configuration Complete!"
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo ""
    log_info "Generated files:"
    log_info "  - $static_nodes_file"
    log_info "  - $permissioned_nodes_file"
    echo ""
    log_info "Next steps:"
    log_info "1. Verify services are running on all containers"
    log_info "2. Check peer connections using curl command"
    log_info "3. Verify discovery settings on RPC nodes - should be disabled for 2503"
}

# Run main function
main "$@"

