#!/usr/bin/env bash
# Collect enodes from all nodes and generate allowlist
# Usage: Update NODES array with your node IPs, then: bash collect-all-enodes.sh

set -euo pipefail

# Load IP configuration
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)"
source "${PROJECT_ROOT}/config/ip-addresses.conf" 2>/dev/null || true


WORK_DIR="${WORK_DIR:-./besu-enodes-$(date +%Y%m%d-%H%M%S)}"
mkdir -p "$WORK_DIR"

# Node inventory: IP:RPC_PORT:USE_RPC (use_rpc=1 if RPC available, 0 for nodekey)
declare -A NODES=(
    ["${IP_SERVICE_13:-${IP_SERVICE_13:-${IP_SERVICE_13:-${IP_SERVICE_13:-${IP_SERVICE_13:-${IP_SERVICE_13:-192.168.11.13}}}}}}"]="8545:1"  # validator-1
    ["${IP_DEVICE_14:-${IP_DEVICE_14:-${IP_DEVICE_14:-${IP_DEVICE_14:-${IP_DEVICE_14:-${IP_DEVICE_14:-192.168.11.14}}}}}}"]="8545:1"  # validator-2
    ["${IP_SERVICE_15:-${IP_SERVICE_15:-192.168.11.15}}"]="8545:1"  # validator-3
    ["${IP_SERVICE_16:-${IP_SERVICE_16:-192.168.11.16}}"]="8545:1"  # validator-4
    ["${IP_SERVICE_18:-${IP_SERVICE_18:-192.168.11.18}}"]="8545:1"  # validator-5
    ["${IP_SERVICE_19:-${IP_SERVICE_19:-192.168.11.19}}"]="8545:1"  # sentry-2
    ["${IP_OMADA:-192.168.11.20}"]="8545:1"  # sentry-3
    ["${IP_SERVICE_21:-${IP_SERVICE_21:-${IP_SERVICE_21:-${IP_SERVICE_21:-${IP_SERVICE_21:-${IP_SERVICE_21:-192.168.11.21}}}}}}"]="8545:1"  # sentry-4
    ["${IP_SERVICE_22:-${IP_SERVICE_22:-192.168.11.22}}"]="8545:1"  # sentry-5
    ["${IP_SERVICE_23:-${IP_SERVICE_23:-192.168.11.23}}"]="8545:1"  # rpc-1
    ["${IP_SERVICE_24:-${IP_SERVICE_24:-192.168.11.24}}"]="8545:1"  # rpc-2
    ["${IP_SERVICE_25:-${IP_SERVICE_25:-192.168.11.25}}"]="8545:1"  # rpc-3
)

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
SSH_USER="${SSH_USER:-root}"
SSH_OPTS="${SSH_OPTS:--o StrictHostKeyChecking=accept-new}"

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"; }

validate_enode() {
    local enode="$1"
    local node_id
    
    node_id=$(echo "$enode" | sed 's|^enode://||' | cut -d'@' -f1 | tr '[:upper:]' '[:lower:]')
    
    if [[ ${#node_id} -ne 128 ]]; then
        return 1
    fi
    
    if ! echo "$node_id" | grep -qE '^[0-9a-f]{128}$'; then
        return 1
    fi
    
    return 0
}

extract_via_rpc() {
    local ip="$1"
    local rpc_port="$2"
    local rpc_url="http://${ip}:${rpc_port}"
    
    local response
    response=$(curl -s -m 5 -X POST \
        -H "Content-Type: application/json" \
        --data '{"jsonrpc":"2.0","method":"admin_nodeInfo","params":[],"id":1}' \
        "${rpc_url}" 2>/dev/null || echo "")
    
    if [[ -z "$response" ]]; then
        return 1
    fi
    
    if echo "$response" | python3 -c "import sys, json; data=json.load(sys.stdin); sys.exit(0 if 'error' not in data else 1)" 2>/dev/null; then
        return 1
    fi
    
    local enode
    enode=$(echo "$response" | python3 -c "import sys, json; print(json.load(sys.stdin).get('result', {}).get('enode', ''))" 2>/dev/null)
    
    if [[ -z "$enode" ]] || [[ "$enode" == "None" ]] || [[ "$enode" == "null" ]]; then
        return 1
    fi
    
    if validate_enode "$enode"; then
        echo "$enode"
        return 0
    fi
    
    return 1
}

extract_via_ssh_nodekey() {
    local ip="$1"
    local ssh_target="${SSH_USER}@${ip}"
    
    local enode
    enode=$(ssh $SSH_OPTS "$ssh_target" bash << REMOTE_SCRIPT
DATA_PATH="/data/besu"
BESU_BIN="/opt/besu/bin/besu"
HOST_IP="${ip}"

for path in "\${DATA_PATH}/key" "\${DATA_PATH}/nodekey" "/keys/besu/nodekey"; do
    if [[ -f "\$path" ]]; then
        ENODE=\$("\${BESU_BIN}" public-key export --node-private-key-file="\$path" --format=enode 2>/dev/null | sed "s/@[0-9.]*:/@\${HOST_IP}:/")
        if [[ -n "\$ENODE" ]]; then
            echo "\$ENODE"
            exit 0
        fi
    fi
done
exit 1
REMOTE_SCRIPT
    )
    
    if [[ -n "$enode" ]] && validate_enode "$enode"; then
        echo "$enode"
        return 0
    fi
    
    return 1
}

log_info "Starting enode collection..."
echo ""

COLLECTED_ENODES="$WORK_DIR/collected-enodes.txt"
DUPLICATES="$WORK_DIR/duplicates.txt"
INVALIDS="$WORK_DIR/invalid-enodes.txt"

> "$COLLECTED_ENODES"
> "$DUPLICATES"
> "$INVALIDS"

declare -A ENODE_BY_IP
declare -A NODE_ID_SET

for ip in "${!NODES[@]}"; do
    IFS=':' read -r rpc_port use_rpc <<< "${NODES[$ip]}"
    
    log_info "Processing node: $ip"
    
    ENODE=""
    
    if [[ "$use_rpc" == "1" ]]; then
        ENODE=$(extract_via_rpc "$ip" "$rpc_port" || echo "")
    fi
    
    if [[ -z "$ENODE" ]]; then
        ENODE=$(extract_via_ssh_nodekey "$ip" || echo "")
    fi
    
    if [[ -z "$ENODE" ]]; then
        log_error "Failed to extract enode from $ip"
        echo "$ip|FAILED" >> "$INVALIDS"
        continue
    fi
    
    if ! validate_enode "$ENODE"; then
        log_error "Invalid enode format from $ip"
        echo "$ip|$ENODE" >> "$INVALIDS"
        continue
    fi
    
    NODE_ID=$(echo "$ENODE" | sed 's|^enode://||' | cut -d'@' -f1)
    ENDPOINT=$(echo "$ENODE" | sed 's|.*@||')
    
    if [[ -n "${NODE_ID_SET[$NODE_ID]:-}" ]]; then
        log_warn "Duplicate node ID: ${NODE_ID:0:32}..."
        echo "$ip|$ENODE|DUPLICATE_NODE_ID|${NODE_ID_SET[$NODE_ID]}" >> "$DUPLICATES"
        continue
    fi
    
    if [[ -n "${ENODE_BY_IP[$ENDPOINT]:-}" ]]; then
        log_warn "Duplicate endpoint: $ENDPOINT"
        echo "$ip|$ENODE|DUPLICATE_ENDPOINT|${ENODE_BY_IP[$ENDPOINT]}" >> "$DUPLICATES"
        continue
    fi
    
    NODE_ID_SET[$NODE_ID]="$ip"
    ENODE_BY_IP[$ENDPOINT]="$ip"
    
    echo "$ip|$ENODE" >> "$COLLECTED_ENODES"
    log_success "Collected: $ip"
done

VALID_COUNT=$(wc -l < "$COLLECTED_ENODES" 2>/dev/null || echo "0")
DUP_COUNT=$(wc -l < "$DUPLICATES" 2>/dev/null || echo "0")
INVALID_COUNT=$(wc -l < "$INVALIDS" 2>/dev/null || echo "0")

echo ""
log_info "Collection Summary:"
log_success "Valid enodes: $VALID_COUNT"
[[ "$DUP_COUNT" -gt 0 ]] && log_warn "Duplicates: $DUP_COUNT (see $DUPLICATES)"
[[ "$INVALID_COUNT" -gt 0 ]] && log_error "Invalid: $INVALID_COUNT (see $INVALIDS)"

echo ""
log_info "Output directory: $WORK_DIR"
log_info "Run: bash ${SCRIPT_DIR}/besu-generate-allowlist.sh $COLLECTED_ENODES"