#!/usr/bin/env bash
# Fix enode URLs using Besu's built-in commands
# This script generates correct enode URLs from actual Besu node keys

set -euo pipefail

# Container IP mapping
declare -A CONTAINER_IPS=(
    [106]="192.168.11.13"  # besu-validator-1
    [107]="192.168.11.14"  # besu-validator-2
    [108]="192.168.11.15"  # besu-validator-3
    [109]="192.168.11.16"  # besu-validator-4
    [110]="192.168.11.18"  # besu-validator-5
    [111]="192.168.11.19"  # besu-sentry-2
    [112]="192.168.11.20"  # besu-sentry-3
    [113]="192.168.11.21"  # besu-sentry-4
    [114]="192.168.11.22"  # besu-sentry-5
    [115]="192.168.11.23"  # besu-rpc-1
    [116]="192.168.11.24"  # besu-rpc-2
    [117]="192.168.11.25"  # besu-rpc-3
)

# 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}[WARNING]${NC} $1"; }
log_error() { echo -e "${RED}[ERROR]${NC} $1"; }

PROXMOX_HOST="${PROXMOX_HOST:-192.168.11.10}"
WORK_DIR="/tmp/besu-enode-fix-$$"
mkdir -p "$WORK_DIR"

cleanup() {
    rm -rf "$WORK_DIR"
}
trap cleanup EXIT

# Extract enode from a running Besu node using admin_nodeInfo
extract_enode_from_running_node() {
    local vmid="$1"
    local ip="${CONTAINER_IPS[$vmid]}"
    
    log_info "Extracting enode from container $vmid ($ip)..."
    
    # Try to get enode via RPC (works if RPC is enabled)
    local enode
    enode=$(ssh -o StrictHostKeyChecking=accept-new "root@${PROXMOX_HOST}" \
        "pct exec $vmid -- curl -s -X POST --data '{\"jsonrpc\":\"2.0\",\"method\":\"admin_nodeInfo\",\"params\":[],\"id\":1}' -H 'Content-Type: application/json' 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" && "$enode" != "null" && "$enode" != "" ]]; then
        echo "$enode"
        return 0
    fi
    
    log_warn "Could not get enode via RPC for container $vmid, trying alternative method..."
    return 1
}

# Extract node key path and generate enode using Besu CLI
extract_enode_from_nodekey() {
    local vmid="$1"
    local ip="${CONTAINER_IPS[$vmid]}"
    
    log_info "Extracting enode from nodekey for container $vmid ($ip)..."
    
    # Find nodekey file (Besu stores it in data directory)
    local nodekey_path
    nodekey_path=$(ssh -o StrictHostKeyChecking=accept-new "root@${PROXMOX_HOST}" \
        "pct exec $vmid -- find /data/besu -name 'nodekey*' -o -name 'key.priv' 2>/dev/null | head -1" || echo "")
    
    if [[ -z "$nodekey_path" ]]; then
        # Try common locations
        for path in "/data/besu/key" "/data/besu/nodekey" "/keys/besu/nodekey"; 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 [[ -z "$nodekey_path" ]]; then
        log_error "Nodekey not found for container $vmid"
        return 1
    fi
    
    log_info "Found nodekey at: $nodekey_path"
    
    # Use Besu to export public key in enode format
    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 container IP
        enode=$(echo "$enode" | sed "s/@[0-9.]*:/@${ip}:/")
        echo "$enode"
        return 0
    fi
    
    log_error "Failed to generate enode from nodekey for container $vmid"
    return 1
}

# Validate enode format
validate_enode() {
    local enode="$1"
    
    # Extract node ID part
    local node_id
    node_id=$(echo "$enode" | sed 's|^enode://||' | cut -d'@' -f1 | tr '[:upper:]' '[:lower:]')
    
    # Check length
    if [[ ${#node_id} -ne 128 ]]; then
        log_error "Invalid enode: node ID length is ${#node_id}, expected 128"
        return 1
    fi
    
    # Check it's valid hex
    if ! echo "$node_id" | grep -qE '^[0-9a-f]{128}$'; then
        log_error "Invalid enode: node ID contains non-hex characters"
        return 1
    fi
    
    log_success "Enode validated: ${node_id:0:16}...${node_id: -16} (128 chars)"
    return 0
}

# Collect all enodes from containers
collect_enodes() {
    log_info "Collecting enodes from all containers..."
    echo "" > "$WORK_DIR/enodes.txt"
    
    local validators=()
    local sentries=()
    local rpcs=()
    
    for vmid in 106 107 108 109 110; do
        local hostname ip enode
        hostname=$(ssh -o StrictHostKeyChecking=accept-new "root@${PROXMOX_HOST}" \
            "pct config $vmid 2>/dev/null | grep '^hostname:' | cut -d' ' -f2" || echo "unknown")
        ip="${CONTAINER_IPS[$vmid]}"
        
        # Try running node first, then nodekey
        enode=$(extract_enode_from_running_node "$vmid" 2>/dev/null || \
                extract_enode_from_nodekey "$vmid" 2>/dev/null || echo "")
        
        if [[ -n "$enode" ]]; then
            if validate_enode "$enode"; then
                echo "$vmid|$hostname|$ip|$enode|validator" >> "$WORK_DIR/enodes.txt"
                validators+=("$enode")
                log_success "Container $vmid ($hostname): enode extracted"
            else
                log_error "Container $vmid: invalid enode format"
            fi
        else
            log_warn "Container $vmid: could not extract enode"
        fi
    done
    
    for vmid in 111 112 113 114; do
        local hostname ip enode
        hostname=$(ssh -o StrictHostKeyChecking=accept-new "root@${PROXMOX_HOST}" \
            "pct config $vmid 2>/dev/null | grep '^hostname:' | cut -d' ' -f2" || echo "unknown")
        ip="${CONTAINER_IPS[$vmid]}"
        
        enode=$(extract_enode_from_running_node "$vmid" 2>/dev/null || \
                extract_enode_from_nodekey "$vmid" 2>/dev/null || echo "")
        
        if [[ -n "$enode" ]]; then
            if validate_enode "$enode"; then
                echo "$vmid|$hostname|$ip|$enode|sentry" >> "$WORK_DIR/enodes.txt"
                sentries+=("$enode")
                log_success "Container $vmid ($hostname): enode extracted"
            fi
        fi
    done
    
    for vmid in 115 116 117; do
        local hostname ip enode
        hostname=$(ssh -o StrictHostKeyChecking=accept-new "root@${PROXMOX_HOST}" \
            "pct config $vmid 2>/dev/null | grep '^hostname:' | cut -d' ' -f2" || echo "unknown")
        ip="${CONTAINER_IPS[$vmid]}"
        
        enode=$(extract_enode_from_running_node "$vmid" 2>/dev/null || \
                extract_enode_from_nodekey "$vmid" 2>/dev/null || echo "")
        
        if [[ -n "$enode" ]]; then
            if validate_enode "$enode"; then
                echo "$vmid|$hostname|$ip|$enode|rpc" >> "$WORK_DIR/enodes.txt"
                rpcs+=("$enode")
                log_success "Container $vmid ($hostname): enode extracted"
            fi
        fi
    done
    
    log_info "Collected ${#validators[@]} validator enodes, ${#sentries[@]} sentry enodes, ${#rpcs[@]} RPC enodes"
}

# Generate corrected static-nodes.json (validators only)
generate_static_nodes_json() {
    log_info "Generating corrected static-nodes.json..."
    
    python3 << PYEOF
import json

# Read collected enodes
enodes = []
with open('$WORK_DIR/enodes.txt', 'r') as f:
    for line in f:
        line = line.strip()
        if not line:
            continue
        parts = line.split('|')
        if len(parts) >= 4 and parts[4] == 'validator':
            enodes.append(parts[3])  # enode URL

# Generate JSON
static_nodes = sorted(set(enodes))  # Remove duplicates, sort

with open('$WORK_DIR/static-nodes.json', 'w') as f:
    json.dump(static_nodes, f, indent=2)

print(f"Generated static-nodes.json with {len(static_nodes)} validators")
PYEOF
    
    log_success "Generated static-nodes.json"
}

# Generate corrected permissions-nodes.toml (all nodes)
generate_permissions_nodes_toml() {
    log_info "Generating corrected permissions-nodes.toml..."
    
    python3 << 'PYEOF'
import re

# Read collected enodes
enodes = []
with open('$WORK_DIR/enodes.txt', 'r') as f:
    for line in f:
        line = line.strip()
        if not line:
            continue
        parts = line.split('|')
        if len(parts) >= 4:
            enodes.append(parts[3])  # enode URL

# Remove duplicates and sort
enodes = sorted(set(enodes))

# Generate TOML
toml_content = """# Node Permissioning Configuration
# Lists nodes that are allowed to connect to this node
# Generated using Besu native commands
# All validators, sentries, and RPC nodes are included

nodes-allowlist=[
"""

for enode in enodes:
    toml_content += f'  "{enode}",\n'

# Remove trailing comma
toml_content = toml_content.rstrip(',\n') + '\n]'

with open('$WORK_DIR/permissions-nodes.toml', 'w') as f:
    f.write(toml_content)

print(f"Generated permissions-nodes.toml with {len(enodes)} nodes")
PYEOF
    
    log_success "Generated permissions-nodes.toml"
}

# Deploy corrected files to containers
deploy_files() {
    log_info "Deploying corrected files to containers..."
    
    # Copy static-nodes.json to all containers
    for vmid in 106 107 108 109 110 111 112 113 114 115 116 117; do
        if ssh -o StrictHostKeyChecking=accept-new "root@${PROXMOX_HOST}" "pct status $vmid 2>/dev/null | grep -q running"; then
            log_info "Deploying to container $vmid..."
            scp -o StrictHostKeyChecking=accept-new "$WORK_DIR/static-nodes.json" \
                "root@${PROXMOX_HOST}:/tmp/static-nodes-${vmid}.json"
            scp -o StrictHostKeyChecking=accept-new "$WORK_DIR/permissions-nodes.toml" \
                "root@${PROXMOX_HOST}:/tmp/permissions-nodes-${vmid}.toml"
            
            ssh -o StrictHostKeyChecking=accept-new "root@${PROXMOX_HOST}" << REMOTE_SCRIPT
pct push $vmid /tmp/static-nodes-${vmid}.json /etc/besu/static-nodes.json
pct push $vmid /tmp/permissions-nodes-${vmid}.toml /etc/besu/permissions-nodes.toml
pct exec $vmid -- chown besu:besu /etc/besu/static-nodes.json /etc/besu/permissions-nodes.toml
rm -f /tmp/static-nodes-${vmid}.json /tmp/permissions-nodes-${vmid}.toml
REMOTE_SCRIPT
            
            log_success "Container $vmid: files deployed"
        else
            log_warn "Container $vmid: not running, skipping"
        fi
    done
}

# Main execution
main() {
    echo "╔════════════════════════════════════════════════════════════════╗"
    echo "║         FIX ENODE URLs USING BESU NATIVE COMMANDS               ║"
    echo "╚════════════════════════════════════════════════════════════════╝"
    echo ""
    
    collect_enodes
    generate_static_nodes_json
    generate_permissions_nodes_toml
    
    echo ""
    log_info "Preview of generated files:"
    echo ""
    echo "=== static-nodes.json ==="
    cat "$WORK_DIR/static-nodes.json" | head -10
    echo ""
    echo "=== permissions-nodes.toml (first 20 lines) ==="
    cat "$WORK_DIR/permissions-nodes.toml" | head -20
    echo ""
    
    read -p "Deploy corrected files to all containers? [y/N]: " -n 1 -r
    echo
    if [[ $REPLY =~ ^[Yy]$ ]]; then
        deploy_files
        log_success "All files deployed successfully!"
        log_info "Next step: Restart Besu services on all containers"
    else
        log_info "Files are in $WORK_DIR (will be cleaned up on exit)"
        log_info "Review them and run deployment manually if needed"
    fi
}

main "$@"