#!/usr/bin/env bash
# Collect enode addresses from all new Besu nodes

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)"
source "${PROJECT_ROOT}/scripts/lib/load-project-env.sh"

declare -A NODE_IPS=(
    [1505]="192.168.11.213"
    [1506]="192.168.11.214"
    [1507]="192.168.11.244"
    [1508]="192.168.11.245"
    [1509]="192.168.11.219"
    [1510]="192.168.11.220"
    [2101]="${RPC_CORE_1:-192.168.11.211}"
    [2102]="192.168.11.212"
    [2103]="192.168.11.217"
    [2201]="${RPC_PUBLIC_1:-192.168.11.221}"
    [2301]="${RPC_PRIVATE_1:-192.168.11.232}"
    [2303]="192.168.11.233"
    [2304]="192.168.11.234"
    [2305]="192.168.11.235"
    [2306]="192.168.11.236"
    [2307]="192.168.11.237"
    [2308]="192.168.11.238"
    [2400]="192.168.11.240"
    [2401]="${RPC_THIRDWEB_1:-192.168.11.241}"
    [2402]="${RPC_THIRDWEB_2:-192.168.11.242}"
    [2403]="${RPC_THIRDWEB_3:-192.168.11.243}"
    [2420]="192.168.11.172"
    [2430]="192.168.11.173"
    [2440]="192.168.11.174"
    [2460]="192.168.11.246"
    [2470]="192.168.11.247"
    [2480]="192.168.11.248"
)

declare -A NODE_NAMES=(
    [1505]="besu-sentry-alltra-1"
    [1506]="besu-sentry-alltra-2"
    [1507]="besu-sentry-hybx-1"
    [1508]="besu-sentry-hybx-2"
    [1509]="besu-sentry-thirdweb-1"
    [1510]="besu-sentry-thirdweb-2"
    [2101]="besu-rpc-core-1"
    [2102]="besu-rpc-core-2"
    [2103]="besu-rpc-admin-core-3"
    [2201]="besu-rpc-public-1"
    [2301]="besu-rpc-private-1"
    [2303]="besu-rpc-private-3"
    [2304]="besu-rpc-private-4"
    [2305]="besu-rpc-private-5"
    [2306]="besu-rpc-private-6"
    [2307]="besu-rpc-private-7"
    [2308]="besu-rpc-private-8"
    [2400]="besu-rpc-thirdweb-primary"
    [2401]="besu-rpc-thirdweb-1"
    [2402]="besu-rpc-thirdweb-2"
    [2403]="besu-rpc-thirdweb-3"
    [2420]="besu-rpc-alltra-1"
    [2430]="besu-rpc-alltra-2"
    [2440]="besu-rpc-alltra-3"
    [2460]="besu-rpc-hybx-1"
    [2470]="besu-rpc-hybx-2"
    [2480]="besu-rpc-hybx-3"
)

TARGET_VMIDS=()

usage() {
    cat <<'EOF'
Usage: ./scripts/collect-all-enodes.sh --vmid <N> [--vmid <N> ...]

Options:
  --vmid <N>   Required. Collect enodes only for the selected VMIDs.
EOF
}

while [[ $# -gt 0 ]]; do
    case "$1" in
        --vmid)
            [[ $# -ge 2 ]] || { usage >&2; exit 2; }
            TARGET_VMIDS+=("$2")
            shift 2
            ;;
        -h|--help)
            usage
            exit 0
            ;;
        *)
            echo "Unknown argument: $1" >&2
            usage >&2
            exit 2
            ;;
    esac
done

[[ ${#TARGET_VMIDS[@]} -gt 0 ]] || { usage >&2; exit 2; }

BLUE='\033[0;34m'
GREEN='\033[0;32m'
NC='\033[0m'

log() { echo -e "${BLUE}[$(date +'%H:%M:%S')]${NC} $1"; }
success() { echo -e "${GREEN}[✓]${NC} $1"; }

collect_enode() {
    local vmid=$1
    local ip=$2
    local hostname=$3
    local host
    host="$(get_host_for_vmid "$vmid")"
    
    log "Collecting enode from $vmid ($hostname)..."
    
    local enode
    enode=$(ssh -o StrictHostKeyChecking=no root@"$host" "pct exec $vmid -- bash -c '
        if [ -f /data/besu/NODE_ID ]; then
            NODE_ID=$(cat /data/besu/NODE_ID)
            echo \"enode://${NODE_ID}@'"$ip"':30303\"
        else
            echo PENDING
        fi
    '" 2>/dev/null || echo "ERROR")
    
    echo "$vmid|$hostname|$ip|$enode"
}

log "==================================="
log "Collecting Enodes from All New Nodes"
log "==================================="
echo ""

ENODE_DIR=$(mktemp -d)

for vmid in "${TARGET_VMIDS[@]}"; do
    if [[ -z "${NODE_IPS[$vmid]:-}" || -z "${NODE_NAMES[$vmid]:-}" ]]; then
        echo "$vmid|unknown|unknown|UNSUPPORTED_VMID" > "$ENODE_DIR/${vmid}.txt"
        continue
    fi
    collect_enode "$vmid" "${NODE_IPS[$vmid]}" "${NODE_NAMES[$vmid]}" > "$ENODE_DIR/${vmid}.txt" &
done

echo ""
log "Waiting for collections..."
wait

echo ""
echo "=========================================="
for file in "$ENODE_DIR"/*.txt; do
    cat "$file"
done

ENODE_LIST="$PROJECT_ROOT/ENODE_COLLECTION_$(date +%Y%m%d_%H%M%S).txt"
{
    echo "# Enode Collection Report"
    echo "# Generated: $(date)"
    echo "# VMID | Hostname | IP | Enode"
    echo "=========================================="
    for file in "$ENODE_DIR"/*.txt; do
        cat "$file"
    done
} > "$ENODE_LIST"

success "Enode collection saved to: $ENODE_LIST"
rm -rf "$ENODE_DIR"