#!/usr/bin/env bash
# Flush mempools on selected Besu nodes (validators, sentries, RPC)
# This script must be run on the Proxmox host.
# Usage:
#   ./flush-all-mempools-proxmox.sh --vmid 2101
#   ./flush-all-mempools-proxmox.sh --apply --vmid 2101

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"

# Colors
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
CYAN='\033[0;36m'
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"; }
log_detail() { echo -e "${CYAN}[DETAIL]${NC} $1"; }

# Check if pct is available
if ! command -v pct &>/dev/null; then
    log_error "This script must be run on the Proxmox host (pct command not found)"
    exit 1
fi

APPLY=false
TARGET_VMIDS=()

usage() {
    cat <<'EOF'
Usage: ./flush-all-mempools-proxmox.sh --vmid <N> [--vmid <N> ...] [--apply]

Options:
  --vmid <N>   Required. Limit restart to one or more VMIDs.
  --apply      Restart services. Without this flag, the script prints the target VMIDs and exits.
EOF
}

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

[[ ${#TARGET_VMIDS[@]} -gt 0 ]] || { log_error "At least one --vmid is required."; usage >&2; exit 2; }

declare -A SERVICE_BY_VMID
for vmid in 1000 1001 1002 1003 1004; do SERVICE_BY_VMID[$vmid]="besu-validator.service"; done
for vmid in 1500 1501 1502 1503 1504 1505 1506 1507 1508 1509 1510; do SERVICE_BY_VMID[$vmid]="besu-sentry.service"; done
for vmid in 2101 2102 2103 2201 2301 2303 2304 2305 2306 2307 2308 2400 2401 2402 2403 2420 2430 2440 2460 2470 2480; do SERVICE_BY_VMID[$vmid]="besu-rpc.service"; done

echo "========================================="
echo "Flush All Besu Mempools"
echo "========================================="
echo ""

if ! $APPLY; then
    log_warn "Dry-run only. Target VMIDs:"
    for vmid in "${TARGET_VMIDS[@]}"; do
        log_info "  VMID $vmid service ${SERVICE_BY_VMID[$vmid]:-unknown}"
    done
    log_info "Re-run with --apply to restart services."
    exit 0
fi

# Function to restart Besu service
restart_besu_service() {
    local vmid=$1
    local service_name=$2
    
    if ! pct status "$vmid" 2>/dev/null | grep -q "running"; then
        log_warn "⚠ VMID $vmid: Container not running - skipping"
        return 1
    fi
    
    log_info "VMID $vmid: Restarting $service_name..."
    
    if pct exec "$vmid" -- systemctl restart "$service_name" 2>/dev/null; then
        log_success "✓ Service restart command sent"
        sleep 2
        return 0
    else
        log_error "✗ Failed to restart service"
        return 1
    fi
}

SUCCESS=0
for vmid in "${TARGET_VMIDS[@]}"; do
    service="${SERVICE_BY_VMID[$vmid]:-besu.service}"
    if restart_besu_service "$vmid" "$service"; then
        ((SUCCESS++))
    fi
done
log_info "Nodes restarted: $SUCCESS/${#TARGET_VMIDS[@]}"
echo ""

# Summary
TOTAL_SUCCESS=$SUCCESS
TOTAL_NODES=${#TARGET_VMIDS[@]}

echo "========================================="
echo "Summary"
echo "========================================="
echo ""
log_success "✓ Successfully restarted: $TOTAL_SUCCESS/$TOTAL_NODES nodes"
echo ""

log_info "Waiting 15 seconds for services to stabilize..."
sleep 15

log_info "Verifying services are running..."
VERIFIED=0
for vmid in "${TARGET_VMIDS[@]}"; do
    if pct status "$vmid" 2>/dev/null | grep -q "running"; then
        if pct exec "$vmid" -- pgrep -f "besu" >/dev/null 2>&1; then
            log_success "✓ VMID $vmid: Besu running"
            ((VERIFIED++))
        else
            log_warn "⚠ VMID $vmid: Besu process not found"
        fi
    fi
done

echo ""
log_success "========================================="
log_success "Mempool Flush Complete!"
log_success "========================================="
log_info ""
log_info "Verified: $VERIFIED/$TOTAL_NODES nodes running Besu"
log_info ""
log_info "Next steps:"
log_info "  1. Wait for all nodes to sync"
log_info "  2. Run: ./scripts/configure-ethereum-mainnet-final.sh"
log_info ""