#!/bin/bash
# os-update.sh — full weekly rolling OS-update cycle.
#
# Behavior:
#   1. preflight cluster health (fail-closed)
#   2. take an etcd snapshot
#   3. compute node order (workers first, control plane last;
#      Stalwart-hosting fsn1 workers moved to end of workers group)
#   4. call update-node.sh for each node, halting on any failure
#   5. finalization: health snapshot + apt history digest
#
# Usage:
#   os-update.sh [--dry-run] [--start-from <node>] [--only <node>]
#
# Environment:
#   CONTROL_PLANE_HOST (default 178.105.17.239)
#   All env vars honored by update-node.sh are honored here as well.
#
# Read OS_UPDATE_PROCEDURE.md alongside this script; the script mirrors it
# step-for-step and the doc is the authoritative reference.
set -euo pipefail

CONTROL_PLANE_HOST="${CONTROL_PLANE_HOST:-178.105.17.239}"
DRY_RUN=0
START_FROM=""
ONLY=""

while [ $# -gt 0 ]; do
  case "$1" in
    --dry-run)     DRY_RUN=1; shift ;;
    --start-from)  START_FROM="$2"; shift 2 ;;
    --only)        ONLY="$2"; shift 2 ;;
    -h|--help)     grep '^#' "$0" | sed 's/^# \{0,1\}//'; exit 0 ;;
    *)             echo "unknown arg: $1" >&2; exit 2 ;;
  esac
done

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
STAMP=$(date -u +%Y%m%dT%H%M%SZ)
LOG_DIR="/tmp/os-update-${STAMP}"
mkdir -p "$LOG_DIR"

log() { echo "[$(date -u +%Y-%m-%dT%H:%M:%SZ)] $*" | tee -a "$LOG_DIR/main.log"; }
die() { log "FATAL: $*"; exit 1; }

log "=== os-update.sh cycle $STAMP ==="
log "log dir: $LOG_DIR"

# --- 0. sanity checks ---------------------------------------------------------
command -v kubectl >/dev/null || die "kubectl not on PATH"
command -v jq >/dev/null      || die "jq not on PATH (needed for health checks)"

# --- 1. preflight -------------------------------------------------------------
log "[preflight] cluster health"
if ! RETRY_ON_TRANSIENT=1 "$SCRIPT_DIR/cluster-health.sh" | tee "$LOG_DIR/preflight.log"; then
  die "cluster is not healthy at preflight — refuse to start OS updates"
fi

# --- 2. etcd snapshot ---------------------------------------------------------
if [ "$DRY_RUN" -eq 0 ]; then
  log "[preflight] taking k3s etcd snapshot"
  ssh -o ConnectTimeout=10 -o StrictHostKeyChecking=accept-new root@"$CONTROL_PLANE_HOST" \
    "k3s etcd-snapshot save --name pre-os-update-$(date +%Y%m%d)" | tee "$LOG_DIR/etcd-snapshot.log" || \
      log "WARN: etcd snapshot failed — continuing but this is a risk on CP reboot"
else
  log "[preflight] DRY-RUN — skipping etcd snapshot"
fi

# --- 3. node ordering ---------------------------------------------------------
# Ordering rule:
#   - workers before control plane
#   - within workers: nodes NOT hosting Stalwart first, Stalwart-hosting fsn1 nodes last
#   - k3s-cp-1 always last
STALWART_NODE=$(kubectl -n stalwart get pod -l app=stalwart -o jsonpath='{.items[*].spec.nodeName}' 2>/dev/null | tr ' ' '\n' | sort -u || true)
# If the pod's not currently up (e.g. Pending) we still want to protect fsn1 workers.
CP_NAME="k3s-cp-1"

ALL_NODES=$(kubectl get nodes -o jsonpath='{.items[*].metadata.name}' | tr ' ' '\n')

workers=()
stalwart_workers=()
for n in $ALL_NODES; do
  [ "$n" = "$CP_NAME" ] && continue
  if [ -n "$STALWART_NODE" ] && [ "$n" = "$STALWART_NODE" ]; then
    stalwart_workers+=("$n")
    continue
  fi
  # Any fsn1 worker is a potential Stalwart host — schedule after non-Stalwart nodes.
  loc=$(kubectl get node "$n" -o jsonpath='{.metadata.labels.csi\.hetzner\.cloud/location}' 2>/dev/null || echo "")
  if [ "$loc" = "fsn1" ]; then
    stalwart_workers+=("$n")
  else
    workers+=("$n")
  fi
done

ORDER=("${workers[@]}" "${stalwart_workers[@]}" "$CP_NAME")

if [ -n "$ONLY" ]; then
  ORDER=("$ONLY")
elif [ -n "$START_FROM" ]; then
  new=()
  skip=1
  for n in "${ORDER[@]}"; do
    [ "$n" = "$START_FROM" ] && skip=0
    [ $skip -eq 0 ] && new+=("$n")
  done
  ORDER=("${new[@]}")
fi

log "[plan] ordered nodes (${#ORDER[@]}): ${ORDER[*]}"
printf '%s\n' "${ORDER[@]}" > "$LOG_DIR/plan.txt"

if [ "$DRY_RUN" -eq 1 ]; then
  log "DRY-RUN — plan written, no node touched. Exiting."
  exit 0
fi

# --- 4. per-node loop ---------------------------------------------------------
updated=()
skipped=()
for n in "${ORDER[@]}"; do
  log "===================================================================="
  log "==> updating $n"
  log "===================================================================="
  NODE_LOG="$LOG_DIR/${n}.log"
  set +e
  "$SCRIPT_DIR/update-node.sh" "$n" 2>&1 | tee "$NODE_LOG"
  rc=${PIPESTATUS[0]}
  set -e
  case $rc in
    0) updated+=("$n") ;;
    3) skipped+=("$n:drain-blocked") ;;
    *) die "update-node.sh failed for $n (rc=$rc). Cycle halted. See $NODE_LOG" ;;
  esac
done

# --- 5. finalization ----------------------------------------------------------
log "===================================================================="
log "==> finalization"
log "===================================================================="

log "[final] cluster health"
"$SCRIPT_DIR/cluster-health.sh" | tee "$LOG_DIR/final-health.log" || \
  die "final cluster health failed after cycle. Do NOT declare success."

log "[final] apt history digest"
{
  for n in "${updated[@]}"; do
    echo "=== $n ==="
    # resolve ssh target via a mini-eval of node_ssh_target-equivalent
    case "$n" in
      k3s-cp-1)           t="root@178.105.17.239" ;;
      k3s-worker-1)       t="-J root@$CONTROL_PLANE_HOST root@10.42.1.2" ;;
      k3s-worker-2)       t="-J root@$CONTROL_PLANE_HOST root@10.42.1.3" ;;
      k3s-worker-3)       t="root@167.233.121.121" ;;
      k3s-worker-4)       t="root@128.140.3.80" ;;
      k3s-worker-5)       t="root@167.233.192.86" ;;
      k3s-update-runner)  t="root@167.233.79.65" ;;
      *)                  echo "  (unknown ssh target)"; continue ;;
    esac
    ssh -o ConnectTimeout=10 -o StrictHostKeyChecking=accept-new $t \
      'zgrep -h "Commandline\|Install\|Upgrade\|Remove" /var/log/apt/history.log* 2>/dev/null | tail -40' 2>/dev/null \
      || echo "  (could not read apt history)"
  done
} | tee "$LOG_DIR/apt-history.log"

log "=== summary ==="
log "updated (${#updated[@]}): ${updated[*]:-none}"
log "skipped (${#skipped[@]}): ${skipped[*]:-none}"
log "logs:   $LOG_DIR"

if [ ${#skipped[@]} -gt 0 ]; then
  log "cycle finished with skipped nodes — return code 4 so the caller can escalate"
  exit 4
fi

log "cycle complete."