#!/usr/bin/env bash
# Upgrade all running Besu containers to the requested version.
# Installs Java 21 where needed, preserves the previous /opt/besu-* directory for rollback,
# and restarts the detected Besu systemd unit in each container.
#
# Usage:
#   bash scripts/upgrade-besu-all-nodes.sh
#   bash scripts/upgrade-besu-all-nodes.sh --dry-run
#   BESU_VERSION=25.12.0 bash scripts/upgrade-besu-all-nodes.sh

set -euo pipefail

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

BESU_VERSION="${BESU_VERSION:-25.12.0}"
BESU_TAR="besu-${BESU_VERSION}.tar.gz"
BESU_DIR="/opt/besu-${BESU_VERSION}"
DOWNLOAD_URL="${BESU_DOWNLOAD_URL:-https://github.com/hyperledger/besu/releases/download/${BESU_VERSION}/${BESU_TAR}}"
JAVA21_FALLBACK_URL="${JAVA21_FALLBACK_URL:-https://api.adoptium.net/v3/binary/latest/21/ga/linux/x64/jre/hotspot/normal/eclipse}"
RPC_HTTP_MAX_ACTIVE_CONNECTIONS="${RPC_HTTP_MAX_ACTIVE_CONNECTIONS:-256}"
RPC_WS_MAX_ACTIVE_CONNECTIONS="${RPC_WS_MAX_ACTIVE_CONNECTIONS:-256}"
LOCAL_CACHE="${LOCAL_CACHE:-/tmp}"
DRY_RUN=false
[[ "${1:-}" == "--dry-run" ]] && DRY_RUN=true

SSH_OPTS=(-o ConnectTimeout=20 -o ServerAliveInterval=15 -o ServerAliveCountMax=3 -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_ok() { echo -e "${GREEN}[OK]${NC} $1"; }
log_warn() { echo -e "${YELLOW}[WARN]${NC} $1"; }
log_err() { echo -e "${RED}[ERROR]${NC} $1"; }

declare -A HOST_BY_VMID
for v in 1000 1001 1002 1500 1501 1502 2101; do HOST_BY_VMID[$v]="${PROXMOX_R630_01:-${PROXMOX_HOST_R630_01:-192.168.11.11}}"; done
for v in 2201 2303 2401; do HOST_BY_VMID[$v]="${PROXMOX_R630_02:-${PROXMOX_HOST_R630_02:-192.168.11.12}}"; done
for v in 1003 1004 1503 1504 1505 1506 1507 1508 2102 2301 2304 2305 2306 2307 2308 2400 2402 2403; do HOST_BY_VMID[$v]="${PROXMOX_ML110:-${PROXMOX_HOST_ML110:-192.168.11.10}}"; done

BESU_VMIDS=(
  1000 1001 1002 1003 1004
  1500 1501 1502 1503 1504 1505 1506 1507 1508
  2101 2102 2201 2301 2303 2304 2305 2306 2307 2308
  2400 2401 2402 2403
)

host_ssh() {
  local host="$1"
  shift
  ssh "${SSH_OPTS[@]}" "root@${host}" "$@"
}

ensure_tarball() {
  local path="${LOCAL_CACHE}/${BESU_TAR}"
  mkdir -p "$LOCAL_CACHE"
  if [[ -f "$path" ]]; then
    log_ok "Using existing $path" >&2
    printf '%s\n' "$path"
    return 0
  fi
  if $DRY_RUN; then
    printf '%s\n' "$path"
    return 0
  fi
  log_info "Downloading ${DOWNLOAD_URL}" >&2
  curl -fsSL -o "$path" "$DOWNLOAD_URL"
  log_ok "Downloaded $path" >&2
  printf '%s\n' "$path"
}

detect_service() {
  local host="$1"
  local vmid="$2"
  host_ssh "$host" "pct exec ${vmid} -- bash -lc 'systemctl list-units --type=service --no-legend 2>/dev/null | awk \"{print \\\$1}\" | grep -iE \"^besu-(validator|sentry|rpc|rpc-core)\\.service$|^besu\\.service$\" | head -1'" 2>/dev/null || true
}

is_running() {
  local host="$1"
  local vmid="$2"
  host_ssh "$host" "pct status ${vmid} 2>/dev/null | awk '{print \$2}'" 2>/dev/null | grep -q '^running$'
}

prepare_host_tarball() {
  local host="$1"
  local local_path="$2"
  local host_tmp="/tmp/${BESU_TAR}"
  if $DRY_RUN; then
    log_info "  [dry-run] would copy ${BESU_TAR} to ${host}:${host_tmp}"
    return 0
  fi
  scp "${SSH_OPTS[@]}" "$local_path" "root@${host}:${host_tmp}" >/dev/null
}

upgrade_node() {
  local host="$1"
  local vmid="$2"
  local service="$3"

  if ! is_running "$host" "$vmid"; then
    log_warn "VMID ${vmid} @ ${host}: not running, skipping"
    return 0
  fi

  if [[ -z "$service" ]]; then
    log_warn "VMID ${vmid} @ ${host}: no Besu service detected, skipping"
    return 0
  fi

  log_info "VMID ${vmid} @ ${host}: upgrading ${service} to Besu ${BESU_VERSION}"

  if $DRY_RUN; then
    log_info "  [dry-run] would install Java 21, extract ${BESU_TAR}, switch /opt/besu, restart ${service}"
    return 0
  fi

  host_ssh "$host" "pct push ${vmid} /tmp/${BESU_TAR} /tmp/${BESU_TAR}" >/dev/null

  host_ssh "$host" "pct exec ${vmid} -- bash -lc '
    set -euo pipefail
    if [[ ! -e /opt/besu ]]; then
      fallback=\$(find /opt -maxdepth 1 -type d -name \"besu-*\" | sort -V | tail -1)
      if [[ -n \"\${fallback:-}\" ]]; then
        ln -sfn \"\$fallback\" /opt/besu
        chown -h besu:besu /opt/besu 2>/dev/null || true
      fi
    elif [[ ! -L /opt/besu ]]; then
      current_semver=\$(/opt/besu/bin/besu --version 2>/dev/null | grep -Eo \"[0-9]+\\.[0-9]+\\.[0-9]+\" | head -1)
      current_version=\"besu-\${current_semver:-}\"
      [[ -z \"\${current_version:-}\" ]] && current_version=besu-backup-pre-${BESU_VERSION}
      if [[ ! -d \"/opt/\${current_version}\" ]]; then
        mv /opt/besu \"/opt/\${current_version}\"
      else
        rm -rf /opt/besu
      fi
      ln -sfn \"/opt/\${current_version}\" /opt/besu
      chown -h besu:besu /opt/besu 2>/dev/null || true
    fi
    java_major=\$(java -version 2>&1 | sed -n \"1s/.*version \\\"\\([0-9][0-9]*\\).*/\\1/p\")
    if [[ -z \"\${java_major:-}\" || \"\$java_major\" -lt 21 ]]; then
      export DEBIAN_FRONTEND=noninteractive
      apt-get update -qq
      apt-get install -y -qq openjdk-21-jre-headless || true
      java_major=\$(java -version 2>&1 | sed -n \"1s/.*version \\\"\\([0-9][0-9]*\\).*/\\1/p\")
      if [[ -z \"\${java_major:-}\" || \"\$java_major\" -lt 21 ]]; then
        command -v curl >/dev/null 2>&1 || apt-get install -y -qq curl ca-certificates
        tmp_jre=/tmp/java21-jre.tar.gz
        curl -fsSL -o \"\$tmp_jre\" '${JAVA21_FALLBACK_URL}'
        tar -tzf \"\$tmp_jre\" > /tmp/java21-jre.list
        extracted_dir=\$(head -1 /tmp/java21-jre.list | cut -d/ -f1)
        rm -f /tmp/java21-jre.list
        tar -xzf \"\$tmp_jre\" -C /opt
        rm -f \"\$tmp_jre\"
        ln -sfn \"/opt/\${extracted_dir}\" /opt/java-21
        update-alternatives --install /usr/bin/java java /opt/java-21/bin/java 2100
      fi
    fi
    config_file=\$(systemctl cat ${service} | sed -n \"s/.*--config-file=\\\\([^ ]*\\\\).*/\\\\1/p\" | tail -1)
    if [[ -n \"\${config_file:-}\" && -f \"\$config_file\" ]]; then
      find /etc/besu -maxdepth 1 -type f -name \"*.toml\" -print0 2>/dev/null | while IFS= read -r -d \"\" toml; do
        sed -i \
          -e \"/^[[:space:]]*miner-enabled[[:space:]]*=.*/d\" \
          -e \"/^[[:space:]]*privacy-enabled[[:space:]]*=.*/d\" \
          \"\$toml\"
        if grep -q \"^rpc-http-enabled=true\" \"\$toml\" && ! grep -q \"^rpc-http-max-active-connections=\" \"\$toml\"; then
          tmp=\$(mktemp)
          awk \"1; /^rpc-http-port=/{print \\\"rpc-http-max-active-connections=${RPC_HTTP_MAX_ACTIVE_CONNECTIONS}\\\"}\" \"\$toml\" > \"\$tmp\"
          cat \"\$tmp\" > \"\$toml\"
          rm -f \"\$tmp\"
        fi
        if grep -q \"^rpc-ws-enabled=true\" \"\$toml\" && ! grep -q \"^rpc-ws-max-active-connections=\" \"\$toml\"; then
          tmp=\$(mktemp)
          awk \"1; /^rpc-ws-port=/{print \\\"rpc-ws-max-active-connections=${RPC_WS_MAX_ACTIVE_CONNECTIONS}\\\"}\" \"\$toml\" > \"\$tmp\"
          cat \"\$tmp\" > \"\$toml\"
          rm -f \"\$tmp\"
        fi
      done
      if ! grep -q \"^data-storage-format=\" \"\$config_file\"; then
        tmp=\$(mktemp)
        awk \"1; /^sync-mode=/{print \\\"data-storage-format=\\\\\\\"FOREST\\\\\\\"\\\"}\" \"\$config_file\" > \"\$tmp\"
        cat \"\$tmp\" > \"\$config_file\"
        rm -f \"\$tmp\"
      fi
    fi
    cd /opt
    if [[ ! -d ${BESU_DIR} ]]; then
      tar -xzf /tmp/${BESU_TAR} -C /opt
    fi
    rm -f /tmp/${BESU_TAR}
    ln -sfn ${BESU_DIR} /opt/besu
    chown -h besu:besu /opt/besu 2>/dev/null || true
    chown -R besu:besu ${BESU_DIR} /opt/besu-* 2>/dev/null || true
    systemctl restart ${service}
  '" || return 1

  local active version
  active=""
  for _ in $(seq 1 24); do
    active="$(host_ssh "$host" "pct exec ${vmid} -- systemctl is-active ${service}" 2>/dev/null || true)"
    [[ "$active" == "active" ]] && break
    sleep 5
  done
  version="$(host_ssh "$host" "pct exec ${vmid} -- bash -lc '/opt/besu/bin/besu --version 2>/dev/null | grep -m1 \"besu/\" || true'" 2>/dev/null || true)"
  if [[ "$active" == "active" ]]; then
    log_ok "  VMID ${vmid}: ${service} active (${version:-version unavailable})"
    return 0
  fi

  log_err "  VMID ${vmid}: ${service} state=${active:-unknown}"
  host_ssh "$host" "pct exec ${vmid} -- journalctl -u ${service} -n 30 --no-pager" 2>/dev/null || true
  return 1
}

log_info "Upgrade Besu fleet to ${BESU_VERSION}"
$DRY_RUN && log_warn "DRY RUN: no changes will be made"
echo

TARBALL_PATH="$(ensure_tarball)"

declare -A VMIDS_ON_HOST
for vmid in "${BESU_VMIDS[@]}"; do
  host="${HOST_BY_VMID[$vmid]:-}"
  [[ -n "$host" ]] || continue
  VMIDS_ON_HOST[$host]+=" ${vmid}"
done

PASS=0
FAIL=0
SKIP=0

for host in "${!VMIDS_ON_HOST[@]}"; do
  log_info "Host ${host}"
  if ! host_ssh "$host" "echo OK" >/dev/null 2>&1; then
    log_err "  Cannot SSH to ${host}"
    ((FAIL++)) || true
    continue
  fi

  prepare_host_tarball "$host" "$TARBALL_PATH"

  for vmid in ${VMIDS_ON_HOST[$host]}; do
    service="$(detect_service "$host" "$vmid")"
    if ! is_running "$host" "$vmid"; then
      log_warn "VMID ${vmid} @ ${host}: not running, skipping"
      ((SKIP++)) || true
      continue
    fi
    if [[ -z "$service" ]]; then
      log_warn "VMID ${vmid} @ ${host}: no Besu unit found, skipping"
      ((SKIP++)) || true
      continue
    fi
    if upgrade_node "$host" "$vmid" "$service"; then
      ((PASS++)) || true
    else
      ((FAIL++)) || true
    fi
    echo
  done

  if ! $DRY_RUN; then
    host_ssh "$host" "rm -f /tmp/${BESU_TAR}" >/dev/null 2>&1 || true
  fi
done

echo "------------------------------------------------------------"
log_info "Upgrade summary: passed=${PASS} skipped=${SKIP} failed=${FAIL}"
echo "------------------------------------------------------------"

if [[ "$FAIL" -gt 0 ]]; then
  exit 1
fi