#!/usr/bin/env bash
set -euo pipefail

# ------------------------------------------------------------
# qortal-devnet-jar-update.sh
# ------------------------------------------------------------

VERSION="1.0.0"

HOST_TEMPLATE="devnet-node-%d.qortal.link"
HOST_FIXED=""
PORT_START=22221
PORT_END=22227

DEFAULT_JAR="${HOME}/qortal/qortal.jar"
JAR_PATH="$DEFAULT_JAR"

REMOTE_DIR="qortal"
REMOTE_JAR="qortal.jar"
REMOTE_USER="${USER}"
REMOTE_HOME="/home/${REMOTE_USER}"

# Stop behavior
STOP_TIMEOUT=180         # total seconds to wait for java to exit
STOP_POLL=2              # poll interval seconds

# Retry behavior (modeled after your botstrap approach)
SSH_RETRIES=4
SSH_DELAY=5              # initial delay (will exponential backoff)
RSYNC_RETRIES=3
RSYNC_DELAY=5

ONLY_PORTS_RAW=()

usage() {
  cat <<EOF
qortal-devnet-jar-update.sh v${VERSION}

Usage: $(basename "$0") [options]

Options:
  --jar PATH              Local qortal.jar path (default: ${DEFAULT_JAR})
  --host HOST             Fixed host for all ports (default: none)
  --host-template STR     Host template for port->node mapping (default: ${HOST_TEMPLATE})
  --ports A-B             Port range (default: ${PORT_START}-${PORT_END})

  --only SPEC             Only run specific ports (repeatable).
                          Examples:
                            --only 22223
                            --only 22221,22224,22227
                            --only 22221-22227

  --remote-dir DIR        Remote qortal dir (default: ${REMOTE_HOME}/${REMOTE_DIR})
  --user USER             SSH user for remote host (default: ${REMOTE_USER})
  --stop-timeout SEC      Wait for stop (default: ${STOP_TIMEOUT})
  --stop-poll SEC         Poll interval (default: ${STOP_POLL})

  --ssh-retries N         SSH retries (default: ${SSH_RETRIES})
  --ssh-delay SEC         SSH initial delay (exp backoff) (default: ${SSH_DELAY})
  --rsync-retries N       rsync retries (default: ${RSYNC_RETRIES})
  --rsync-delay SEC       rsync delay between retries (default: ${RSYNC_DELAY})

  -h, --help              Help

Examples:
  ./$(basename "$0")
  ./$(basename "$0") --jar /home/crowetic/build/qortal.jar
  ./$(basename "$0") --only 22223
  ./$(basename "$0") --only 22223,22225
  ./$(basename "$0") --only 22225-22227
EOF
}

log() { printf '[%s] %s\n' "$(date +'%F %T')" "$*" >&2; }
warn() { log "WARNING: $*"; }

parse_ports() {
  local s="$1"
  [[ "$s" =~ ^([0-9]+)-([0-9]+)$ ]] || { echo "Invalid --ports (use A-B)" >&2; exit 1; }
  PORT_START="${BASH_REMATCH[1]}"
  PORT_END="${BASH_REMATCH[2]}"
  (( PORT_START <= PORT_END )) || { echo "--ports start must be <= end" >&2; exit 1; }
}

parse_ports_range() {
  local s="$1"
  [[ "$s" =~ ^([0-9]+)-([0-9]+)$ ]] || { echo "Invalid range: $s" >&2; exit 1; }
  local a="${BASH_REMATCH[1]}" b="${BASH_REMATCH[2]}"
  (( a <= b )) || { echo "Range start must be <= end: $s" >&2; exit 1; }
  for ((p=a; p<=b; p++)); do printf '%s\n' "$p"; done
}

while [[ $# -gt 0 ]]; do
  case "$1" in
    --jar) JAR_PATH="${2:-}"; shift 2;;
    --host) HOST_FIXED="${2:-}"; shift 2;;
    --host-template) HOST_TEMPLATE="${2:-}"; shift 2;;
    --ports) parse_ports "${2:-}"; shift 2;;
    --only) ONLY_PORTS_RAW+=( "${2:-}" ); shift 2;;
    --remote-dir) REMOTE_DIR="${2:-}"; shift 2;;
    --user) REMOTE_USER="${2:-}"; shift 2;;
    --stop-timeout) STOP_TIMEOUT="${2:-}"; shift 2;;
    --stop-poll) STOP_POLL="${2:-}"; shift 2;;

    --ssh-retries) SSH_RETRIES="${2:-}"; shift 2;;
    --ssh-delay) SSH_DELAY="${2:-}"; shift 2;;
    --rsync-retries) RSYNC_RETRIES="${2:-}"; shift 2;;
    --rsync-delay) RSYNC_DELAY="${2:-}"; shift 2;;

    -h|--help) usage; exit 0;;
    *) echo "Unknown arg: $1" >&2; usage; exit 1;;
  esac
done

[[ -f "$JAR_PATH" ]] || { echo "Local jar not found: $JAR_PATH" >&2; exit 1; }
[[ "$STOP_TIMEOUT" =~ ^[0-9]+$ ]] || { echo "--stop-timeout must be integer seconds" >&2; exit 1; }
[[ "$STOP_POLL" =~ ^[0-9]+$ ]] || { echo "--stop-poll must be integer seconds" >&2; exit 1; }
[[ "$SSH_RETRIES" =~ ^[0-9]+$ ]] || { echo "--ssh-retries must be integer" >&2; exit 1; }
[[ "$SSH_DELAY" =~ ^[0-9]+$ ]] || { echo "--ssh-delay must be integer seconds" >&2; exit 1; }
[[ "$RSYNC_RETRIES" =~ ^[0-9]+$ ]] || { echo "--rsync-retries must be integer" >&2; exit 1; }
[[ "$RSYNC_DELAY" =~ ^[0-9]+$ ]] || { echo "--rsync-delay must be integer seconds" >&2; exit 1; }

# Build port list
PORT_LIST=()
if (( ${#ONLY_PORTS_RAW[@]} > 0 )); then
  tmp_ports=()
  for spec in "${ONLY_PORTS_RAW[@]}"; do
    if [[ "$spec" =~ ^[0-9]+$ ]]; then
      tmp_ports+=( "$spec" )
    elif [[ "$spec" =~ ^[0-9]+-[0-9]+$ ]]; then
      while IFS= read -r p; do tmp_ports+=( "$p" ); done < <(parse_ports_range "$spec")
    else
      IFS=',' read -r -a parts <<< "$spec"
      for part in "${parts[@]}"; do
        part="${part//[[:space:]]/}"
        [[ -z "$part" ]] && continue
        if [[ "$part" =~ ^[0-9]+$ ]]; then
          tmp_ports+=( "$part" )
        elif [[ "$part" =~ ^[0-9]+-[0-9]+$ ]]; then
          while IFS= read -r p; do tmp_ports+=( "$p" ); done < <(parse_ports_range "$part")
        else
          echo "Invalid --only spec segment: '$part'" >&2
          exit 1
        fi
      done
    fi
  done
  mapfile -t PORT_LIST < <(printf '%s\n' "${tmp_ports[@]}" | awk 'NF' | sort -n | uniq)
else
  for ((p=PORT_START; p<=PORT_END; p++)); do PORT_LIST+=( "$p" ); done
fi

LOCAL_SHA="$(sha256sum "$JAR_PATH" | awk '{print $1}')"

REMOTE_HOME="/home/${REMOTE_USER}"
log "Local jar: $JAR_PATH"
log "Local sha256: $LOCAL_SHA"
if [[ -n "$HOST_FIXED" ]]; then
  log "Target host: ${HOST_FIXED}"
else
  log "Host template: ${HOST_TEMPLATE}"
fi
log "Ports: ${PORT_LIST[*]}"
log "Remote path: ${REMOTE_HOME}/${REMOTE_DIR}/${REMOTE_JAR}"
log "Remote user: ${REMOTE_USER}"
log "SSH retry: ${SSH_RETRIES} attempts, initial delay ${SSH_DELAY}s (exp)"
log "rsync retry: ${RSYNC_RETRIES} attempts, delay ${RSYNC_DELAY}s"

# -------------------------
# Retry helpers (botstrap-style)
# -------------------------

host_for_port() {
  local port="$1"
  if [[ -n "$HOST_FIXED" ]]; then
    echo "$HOST_FIXED"
    return 0
  fi
  local idx=$((port - PORT_START + 1))
  if (( idx < 1 )); then
    echo "Invalid port ${port} for template mapping (PORT_START=${PORT_START})" >&2
    return 1
  fi
  printf "$HOST_TEMPLATE" "$idx"
}

ssh_with_retry_port() {
  local port="$1"
  local host="$2"
  local cmd="$3"
  local retries="$SSH_RETRIES"
  local delay="$SSH_DELAY"

  for ((i=1; i<=retries; i++)); do
    local target="${REMOTE_USER}@${host}"
    >&2 echo "SSH attempt $i/$retries to ${target}:${port} ..."
    >&2 echo "SSH cmd: ssh -p${port} ${target} \"${cmd}\""
    local output=""
    if output=$(
      ssh -p${port} "${target}" "$cmd"
    ); then
      echo "$output"
      return 0
    fi

    >&2 echo "SSH failed. Retrying in $delay seconds..."
    sleep "$delay"
    delay=$((delay * 2))
  done

  >&2 echo "SSH to ${REMOTE_USER}@${host}:${port} failed after $retries attempts."
  return 1
}

rsync_with_retry_port() {
  local port="$1"
  local src="$2"
  local dest="$3" # remote like host:dir/file

  local attempt=1
  local delay="$RSYNC_DELAY"

  # Use the minimal SSH form for rsync's transport.
  local rsh="ssh -p${port}"
  echo "rsync cmd: rsync -raPz -e 'ssh -p${port}' ${src} ${dest}"

  until rsync -raPz -e "ssh -p${port}" ${src} ${dest}; do
    echo "⚠️  rsync failed (attempt $attempt/$RSYNC_RETRIES): $src → $dest"
    if (( attempt >= RSYNC_RETRIES )); then
      echo "❌ Giving up after $RSYNC_RETRIES failed attempts."
      return 1
    fi
    ((attempt++))
    sleep "$delay"
    # you *can* exponential this too if you want; keeping it simple like your rsync wrapper
  done
  echo "✅ rsync succeeded: $src → $dest"
  return 0
}

# -------------------------
# Remote operations
# -------------------------

remote_is_running() {
  local port="$1"
  local host="$2"
  # Return 0 if running, 1 if not running, 2 if unable to determine
  local out
  log "[${port}] run check: ssh -p${port} ${REMOTE_USER}@${host} \"pgrep -fa '[j]ava.*${REMOTE_JAR}' || true\""
  if ! out="$(ssh_with_retry_port "$port" "$host" "pgrep -fa '[j]ava.*${REMOTE_JAR}' || true")"; then
    return 2
  fi
  if [[ -n "$out" ]]; then
    log "[${port}] run check output: ${out}"
  else
    log "[${port}] run check output: (none)"
  fi
  if [[ -n "$out" ]]; then return 0; else return 1; fi
}

remote_stop_wait_or_kill() {
  local port="$1"
  local host="$2"

  remote_is_running "$port" "$host"
  rc=$?
  if [[ "$rc" == "1" ]]; then
    log "[${port}] already stopped ✅"
    return 0
  fi

  log "[${port}] running ./stop.sh ..."
  # stop.sh should block; if it fails, we continue into wait/kill logic.
  if ! ssh_with_retry_port "$port" "$host" "cd ${REMOTE_HOME}/${REMOTE_DIR} && test -x ./stop.sh && ./stop.sh"; then
    remote_is_running "$port" "$host"
    rc=$?
    if [[ "$rc" == "1" ]]; then
      log "[${port}] stop.sh failed but process already stopped ✅"
      return 0
    fi
    warn "[${port}] stop.sh missing/failed; will enforce stop via wait/kill."
  fi

  log "[${port}] waiting up to ${STOP_TIMEOUT}s for java.*${REMOTE_JAR} to exit..."
  local waited=0
  while true; do
    local state
    if remote_is_running "$port" "$host"; then
      state="running"
    else
      rc=$?
      if [[ "$rc" == "1" ]]; then
        state="stopped"
      else
        # couldn't determine due to SSH issues; keep waiting (but counted)
        state="unknown"
      fi
    fi

    if [[ "$state" == "stopped" ]]; then
      log "[${port}] stopped ✅"
      return 0
    fi

    if (( waited >= STOP_TIMEOUT )); then
      warn "[${port}] stop timeout reached (state=${state}) -> killing java.*${REMOTE_JAR}"

      ssh_with_retry_port "$port" "$host" "sudo pkill -f \"java.*${REMOTE_JAR}\" || pkill -f \"java.*${REMOTE_JAR}\" || true" || true
      sleep 2
      # Last check
      if remote_is_running "$port" "$host"; then
        ssh_with_retry_port "$port" "$host" "sudo pkill -9 -f \"java.*${REMOTE_JAR}\" || pkill -9 -f \"java.*${REMOTE_JAR}\" || true" || true
      fi

      # Final verification
      if remote_is_running "$port" "$host"; then
        warn "[${port}] still running after kill attempts."
        return 1
      fi

      log "[${port}] killed -> stopped ✅"
      return 0
    fi

    sleep "$STOP_POLL"
    waited=$((waited + STOP_POLL))
  done
}

remote_backup_jar() {
  local port="$1"
  local host="$2"
  log "[${port}] backing up existing jar (if present)..."
  ssh_with_retry_port "$port" "$host" "cd ${REMOTE_HOME}/${REMOTE_DIR} && \
    if [[ -f \"${REMOTE_JAR}\" ]]; then \
      ts=\$(date +%Y%m%d-%H%M%S); \
      cp -a \"${REMOTE_JAR}\" \"backup-\${ts}-${REMOTE_JAR}\"; \
      echo \"backup-\${ts}-${REMOTE_JAR}\"; \
    else \
      echo \"(no existing jar)\"; \
    fi"
}

remote_verify_jar() {
  local port="$1"
  local host="$2"
  ssh_with_retry_port "$port" "$host" "cd ${REMOTE_HOME}/${REMOTE_DIR} && sha256sum ${REMOTE_JAR} | awk '{print \$1}'"
}

remote_start() {
  local port="$1"
  local host="$2"
  log "[${port}] starting via ./start.sh ..."
  ssh_with_retry_port "$port" "$host" "cd ${REMOTE_HOME}/${REMOTE_DIR} && test -x ./start.sh && nohup ./start.sh >/tmp/qortal-start.log 2>&1 &"
}

# -------------------------
# Main loop
# -------------------------

for port in "${PORT_LIST[@]}"; do
  host="$(host_for_port "$port")" || { warn "[${port}] invalid host mapping; skipping."; continue; }
  log "==================== port ${port} ===================="
  log "[${port}] host: ${host}"

  # 1) stop (wait, kill if needed)
  if ! remote_stop_wait_or_kill "$port" "$host"; then
    warn "[${port}] stop failed; skipping this node."
    continue
  fi

  # 2) backup
  if ! remote_backup_jar "$port" "$host" >/dev/null; then
    warn "[${port}] backup step failed (continuing anyway)."
  fi

  # 3) rsync jar (with retries)
  log "[${port}] rsync new jar..."
  if ! rsync_with_retry_port "$port" "$JAR_PATH" "${REMOTE_USER}@${host}:${REMOTE_HOME}/${REMOTE_DIR}/${REMOTE_JAR}"; then
    warn "[${port}] rsync failed; skipping start."
    continue
  fi

  # 4) verify checksum
  log "[${port}] verifying sha256..."
  remote_sha="$(remote_verify_jar "$port" "$host" 2>/dev/null | tail -n 1 | tr -d '\r\n' || true)"
  if [[ -z "$remote_sha" ]]; then
    warn "[${port}] could not read remote sha; continuing."
  elif [[ "$remote_sha" != "$LOCAL_SHA" ]]; then
    warn "[${port}] SHA MISMATCH remote=$remote_sha local=$LOCAL_SHA"
  else
    log "[${port}] sha match ✅"
  fi

  # 5) start
  if ! remote_start "$port" "$host"; then
    warn "[${port}] start failed."
    continue
  fi

  log "[${port}] done."
done

log "All nodes processed."