#!/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 (non-Stalwart first, Stalwart-hosting fsn1 workers last)
#        then control-plane nodes (etcd followers first, leader last).
#   3b. SSH-key drift preflight: verify operator can reach every planned node
#        and that cp-1's pubkey is present in root's authorized_keys (auto-heal
#        by appending if missing; hard-fail if unreachable). See DEV-587.
#   4. call update-node.sh for each worker and update-cp-node.sh for each CP,
#      halting on any failure
#   5. finalization: health snapshot + apt history digest
#
# Usage:
#   os-update.sh [--dry-run] [--start-from <node>] [--only <node>]
#                [--skip-cps] [--only-cps] [--skip-ssh-preflight]
#
# Flags:
#   --skip-cps            run workers only (legacy behavior, useful for narrow patches).
#   --only-cps            skip workers, only walk the CP phase.
#   --only <n>            update a single node (worker or CP, routed automatically).
#   --skip-ssh-preflight  skip the SSH-drift preflight (escape hatch for bootstrap
#                         scenarios where cp-1's pubkey is intentionally absent).
#
# Environment:
#   CONTROL_PLANE_HOST (default 178.105.17.239)
#   POST_CP_WAIT_SECONDS (default 180) — settle between CPs after uncordon.
#   All env vars honored by update-node.sh / update-cp-node.sh are honored here.
#
# Read OS_UPDATE_PROCEDURE.md alongside this script; the script mirrors it
# step-for-step and the doc is the authoritative reference. CP handling is
# delegated to update-cp-node.sh; see CP_UPDATE_PROCEDURE.md.
set -euo pipefail

CONTROL_PLANE_HOST="${CONTROL_PLANE_HOST:-178.105.17.239}"
POST_CP_WAIT_SECONDS="${POST_CP_WAIT_SECONDS:-180}"
DRY_RUN=0
START_FROM=""
ONLY=""
SKIP_CPS=0
ONLY_CPS=0
SKIP_SSH_PREFLIGHT="${SKIP_SSH_PREFLIGHT:-0}"

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 ;;
    --skip-cps)            SKIP_CPS=1; shift ;;
    --only-cps)            ONLY_CPS=1; shift ;;
    --skip-ssh-preflight)  SKIP_SSH_PREFLIGHT=1; shift ;;
    -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; }

# node_ssh_endpoint <node-name>
# Prints the ssh argv fragment (may include `-J bastion`) used to reach a node.
# Private-IP workers go through cp-1 via ProxyJump; every other node has a
# public IP. Returns non-zero for an unknown node. Callers pass the output
# unquoted so `-J` is treated as its own arg.
node_ssh_endpoint() {
  case "$1" in
    k3s-cp-1)     echo "root@178.105.17.239" ;;
    k3s-cp-2)     echo "root@188.245.85.199" ;;
    k3s-cp-3)     echo "root@49.13.92.162" ;;
    k3s-worker-1) echo "-J root@$CONTROL_PLANE_HOST root@10.42.1.2" ;;
    k3s-worker-2) echo "-J root@$CONTROL_PLANE_HOST root@10.42.1.3" ;;
    k3s-worker-3) echo "root@167.233.121.121" ;;
    k3s-worker-4) echo "root@128.140.3.80" ;;
    k3s-worker-5) echo "root@167.233.192.86" ;;
    k3s-worker-6) echo "root@49.12.223.133" ;;
    k3s-update-runner) echo "root@167.233.79.65" ;;
    *)            return 1 ;;
  esac
}

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 first (non-Stalwart, then Stalwart-hosting fsn1 last)
#   - control-plane nodes AFTER workers, one at a time, etcd followers first,
#     leader last. Rationale (see CP_UPDATE_PROCEDURE.md):
#     * two of three CPs down = etcd quorum loss = write outage. Only one CP
#       cordoned/rebooting at a time.
#     * draining a follower is a no-op for raft leadership; draining the leader
#       forces a re-election, so doing followers first minimises leader flapping.
#   Since DEV-510 (2026-08-22) the cluster runs HA embedded etcd across
#   cp-1/cp-2/cp-3 with 3-node quorum. This script delegates all CP steps to
#   update-cp-node.sh, which handles swap (idempotent), preflight, drain, apt,
#   reboot with /livez polling, etcd rejoin check, and finalize.
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_NODES=$(kubectl get nodes -l node-role.kubernetes.io/control-plane -o jsonpath='{.items[*].metadata.name}' 2>/dev/null | tr ' ' '\n' || true)
ALL_NODES=$(kubectl get nodes -o jsonpath='{.items[*].metadata.name}' | tr ' ' '\n')

is_cp() {
  local n="$1"
  for c in $CP_NODES; do [ "$c" = "$n" ] && return 0; done
  return 1
}

workers=()
stalwart_workers=()
for n in $ALL_NODES; do
  is_cp "$n" && 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

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

# --- 3b. control-plane ordering (leader last) --------------------------------
# CP-node → public IP map. Kept out-of-cluster so a full-cluster incident
# doesn't hide etcd from the operator machine.
declare -A CP_HOST=(
  [k3s-cp-1]="178.105.17.239"
  [k3s-cp-2]="188.245.85.199"
  [k3s-cp-3]="49.13.92.162"
)

etcd_leader_ip() {
  # Query etcd leader endpoint from the operator's control-plane host. Uses
  # the same TLS material as update-cp-node.sh. Prints "IP" on success,
  # nothing on failure.
  ssh -o ConnectTimeout=10 -o StrictHostKeyChecking=accept-new -o BatchMode=yes \
    root@"$CONTROL_PLANE_HOST" \
    'ETCDCTL_API=3 etcdctl \
      --endpoints=https://127.0.0.1:2379 \
      --cacert=/var/lib/rancher/k3s/server/tls/etcd/server-ca.crt \
      --cert=/var/lib/rancher/k3s/server/tls/etcd/server-client.crt \
      --key=/var/lib/rancher/k3s/server/tls/etcd/server-client.key \
      endpoint status --cluster -w simple 2>/dev/null' \
    | awk -F, '$5 ~ /true/ {print $1}' \
    | sed -E 's|https?://([^:]+):.*|\1|'
}

CP_ORDER=()
if [ -n "$CP_NODES" ]; then
  LEADER_IP="$(etcd_leader_ip || true)"
  leader_node=""
  followers=()
  for n in $CP_NODES; do
    ip="${CP_HOST[$n]:-}"
    if [ -n "$ip" ] && [ "$ip" = "$LEADER_IP" ]; then
      leader_node="$n"
    else
      followers+=("$n")
    fi
  done
  if [ -n "$leader_node" ]; then
    CP_ORDER=("${followers[@]}" "$leader_node")
    log "[plan] etcd leader is $leader_node ($LEADER_IP) — updating followers first, leader last"
  else
    # Couldn't identify leader — refuse to guess. Fall back to no CPs.
    CP_ORDER=()
    log "[plan] WARN could not identify etcd leader — skipping CP phase this cycle"
  fi
fi

# Route --only <node> automatically to worker or CP list.
if [ -n "$ONLY" ]; then
  if is_cp "$ONLY"; then
    CP_ORDER=("$ONLY")
    WORKER_ORDER=()
  else
    WORKER_ORDER=("$ONLY")
    CP_ORDER=()
  fi
elif [ -n "$START_FROM" ]; then
  new=()
  skip=1
  for n in "${WORKER_ORDER[@]}"; do
    [ "$n" = "$START_FROM" ] && skip=0
    [ $skip -eq 0 ] && new+=("$n")
  done
  WORKER_ORDER=("${new[@]}")
  # If START_FROM is a CP, trim CP_ORDER instead.
  if is_cp "$START_FROM"; then
    new_cp=()
    skip=1
    for n in "${CP_ORDER[@]:-}"; do
      [ -z "$n" ] && continue
      [ "$n" = "$START_FROM" ] && skip=0
      [ $skip -eq 0 ] && new_cp+=("$n")
    done
    CP_ORDER=("${new_cp[@]}")
    WORKER_ORDER=()
  fi
fi

[ "$SKIP_CPS" -eq 1 ] && CP_ORDER=()
[ "$ONLY_CPS" -eq 1 ] && WORKER_ORDER=()

log "[plan] workers (${#WORKER_ORDER[@]}): ${WORKER_ORDER[*]:-none}"
log "[plan] control-planes (${#CP_ORDER[@]}): ${CP_ORDER[*]:-none}"
{
  echo "# workers"
  printf '%s\n' "${WORKER_ORDER[@]:-}"
  echo "# control-planes"
  printf '%s\n' "${CP_ORDER[@]:-}"
} > "$LOG_DIR/plan.txt"

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

# --- 3c. SSH-key drift preflight ---------------------------------------------
# DEV-587 / DEV-586 (2026-09-27) — some k3s nodes drift to lose cp-1's SSH
# pubkey from /root/.ssh/authorized_keys (probable Hetzner cloud-init resync
# from the HC SSH-keys API; the operator key is NOT in that keyset). When this
# happens mid-cycle the whole run aborts partway through drain / peer-probe.
# This preflight proves the operator can SSH to every planned node and that
# cp-1's pubkey is present in root's authorized_keys — auto-heals drift via
# the operator's own SSH access, and hard-fails BEFORE anything is drained if
# a node is genuinely unreachable.
ssh_drift_preflight() {
  local ok=() repaired=() problem=()

  log "[preflight] SSH-key drift check for planned nodes"

  local cp1_pubkey cp1_key_body
  if ! cp1_pubkey=$(ssh -o BatchMode=yes -o ConnectTimeout=5 -o StrictHostKeyChecking=accept-new \
                         root@"$CONTROL_PLANE_HOST" 'cat /root/.ssh/id_rsa.pub' 2>/dev/null); then
    die "SSH-drift preflight: cannot read cp-1 pubkey from $CONTROL_PLANE_HOST — fix operator->cp-1 SSH first"
  fi
  cp1_pubkey=$(printf '%s' "$cp1_pubkey" | tr -d '\r\n')
  cp1_key_body=$(printf '%s' "$cp1_pubkey" | awk '{print $2}')
  [ -n "$cp1_key_body" ] || die "SSH-drift preflight: could not parse cp-1 pubkey (got: $cp1_pubkey)"
  log "  cp-1 pubkey OK (${cp1_key_body:0:32}...)"

  local plan_nodes=()
  for n in "${WORKER_ORDER[@]:-}" "${CP_ORDER[@]:-}"; do
    [ -z "$n" ] && continue
    plan_nodes+=("$n")
  done
  if [ ${#plan_nodes[@]} -eq 0 ]; then
    log "  (no planned nodes — nothing to preflight)"
    return 0
  fi

  for n in "${plan_nodes[@]}"; do
    local endpoint
    if ! endpoint=$(node_ssh_endpoint "$n"); then
      log "  $n: unknown ssh endpoint — refusing"
      problem+=("$n (no ssh mapping in node_ssh_endpoint)")
      continue
    fi

    # 1. reachability from the operator machine.
    # shellcheck disable=SC2086  # $endpoint may include `-J bastion`; deliberate word-split.
    if ! ssh -o BatchMode=yes -o ConnectTimeout=5 -o StrictHostKeyChecking=accept-new \
             $endpoint true 2>/dev/null; then
      problem+=("$n ($endpoint - operator SSH refused)")
      continue
    fi

    # 2. cp-1 is the source of the key; drift check would be pointless on it.
    if [ "$n" = "k3s-cp-1" ]; then
      ok+=("$n")
      continue
    fi

    # 3. cp-1 pubkey present?
    # shellcheck disable=SC2086
    if ssh -o BatchMode=yes -o ConnectTimeout=5 -o StrictHostKeyChecking=accept-new \
           $endpoint "grep -qF '$cp1_key_body' /root/.ssh/authorized_keys" 2>/dev/null; then
      ok+=("$n")
      continue
    fi

    # 4. auto-heal — append cp-1's pubkey via the (already verified) operator SSH.
    log "  $n: cp-1 pubkey MISSING from /root/.ssh/authorized_keys — appending"
    # shellcheck disable=SC2086
    if ssh -o BatchMode=yes -o ConnectTimeout=5 -o StrictHostKeyChecking=accept-new \
           $endpoint "install -d -m 700 /root/.ssh && printf '%s\n' '$cp1_pubkey' >> /root/.ssh/authorized_keys && chmod 600 /root/.ssh/authorized_keys" \
           2>/dev/null; then
      # shellcheck disable=SC2086
      if ssh -o BatchMode=yes -o ConnectTimeout=5 -o StrictHostKeyChecking=accept-new \
             $endpoint "grep -qF '$cp1_key_body' /root/.ssh/authorized_keys" 2>/dev/null; then
        repaired+=("$n")
      else
        problem+=("$n (auto-heal wrote but grep-verify failed)")
      fi
    else
      problem+=("$n ($endpoint - auto-heal write refused)")
    fi
  done

  log "[preflight] SSH-drift summary:"
  log "  ok        (${#ok[@]}): ${ok[*]:-none}"
  log "  repaired  (${#repaired[@]}): ${repaired[*]:-none}"
  log "  problem   (${#problem[@]}): ${problem[*]:-none}"

  {
    echo "# SSH-drift preflight — $(date -u +%FT%TZ)"
    echo "# cp-1 pubkey: $cp1_pubkey"
    printf 'ok\t%s\n'        "${ok[*]:-}"
    printf 'repaired\t%s\n'  "${repaired[*]:-}"
    printf 'problem\t%s\n'   "${problem[*]:-}"
  } > "$LOG_DIR/ssh-preflight.txt"

  if [ ${#problem[@]} -gt 0 ]; then
    log "SSH-drift preflight FAILED for: ${problem[*]}"
    log "  cp-1 pubkey to install:"
    log "    $cp1_pubkey"
    log "  From any machine that CAN reach the failing node, run:"
    log "    ssh root@<node> \"grep -qF '$cp1_key_body' /root/.ssh/authorized_keys || echo '$cp1_pubkey' >> /root/.ssh/authorized_keys\""
    die "SSH-drift preflight failed — refuse to start rolling update"
  fi
}

if [ "$SKIP_SSH_PREFLIGHT" -eq 1 ]; then
  log "[preflight] SSH-drift check SKIPPED (--skip-ssh-preflight)"
else
  ssh_drift_preflight
fi

# --- 4. worker loop ----------------------------------------------------------
updated=()
skipped=()
for n in "${WORKER_ORDER[@]}"; do
  [ -z "$n" ] && continue
  log "===================================================================="
  log "==> updating worker $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

# --- 4b. control-plane loop --------------------------------------------------
# One CP at a time, followers first, leader last. Delegates to update-cp-node.sh
# --run (Phase A add-swap is idempotent, Phase B..F run with ASSUME_YES=1).
# Between CPs we run cluster-health.sh and re-verify etcd cluster status; halt
# on any yellow signal.
if [ ${#CP_ORDER[@]} -gt 0 ]; then
  log "===================================================================="
  log "==> control-plane phase (${#CP_ORDER[@]} nodes): ${CP_ORDER[*]}"
  log "===================================================================="
  first_cp=1
  for cp in "${CP_ORDER[@]}"; do
    [ -z "$cp" ] && continue
    if [ $first_cp -eq 0 ]; then
      log "-- health gate between CPs (settle ${POST_CP_WAIT_SECONDS}s, then re-check)"
      sleep "$POST_CP_WAIT_SECONDS"
      if ! RETRY_ON_TRANSIENT=1 "$SCRIPT_DIR/cluster-health.sh" 2>&1 | tee "$LOG_DIR/inter-cp-health-before-${cp}.log"; then
        die "cluster health degraded between CPs — halting before $cp. See $LOG_DIR/"
      fi
      # Re-check etcd cluster status; every member must still be started.
      ssh -o ConnectTimeout=10 -o StrictHostKeyChecking=accept-new root@"$CONTROL_PLANE_HOST" \
        'ETCDCTL_API=3 etcdctl \
          --endpoints=https://127.0.0.1:2379 \
          --cacert=/var/lib/rancher/k3s/server/tls/etcd/server-ca.crt \
          --cert=/var/lib/rancher/k3s/server/tls/etcd/server-client.crt \
          --key=/var/lib/rancher/k3s/server/tls/etcd/server-client.key \
          endpoint status --cluster -w table' 2>&1 | tee -a "$LOG_DIR/inter-cp-health-before-${cp}.log" || \
        die "etcd endpoint status failed between CPs — halting before $cp"
    fi
    log "===================================================================="
    log "==> updating control-plane $cp"
    log "===================================================================="
    CP_LOG="$LOG_DIR/${cp}.log"
    set +e
    ASSUME_YES=1 "$SCRIPT_DIR/update-cp-node.sh" "$cp" --run 2>&1 | tee "$CP_LOG"
    rc=${PIPESTATUS[0]}
    set -e
    case $rc in
      0) updated+=("$cp") ;;
      *) die "update-cp-node.sh failed for $cp (rc=$rc). Cycle halted. See $CP_LOG" ;;
    esac
    first_cp=0
  done
fi

# --- 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 ==="
    if ! t=$(node_ssh_endpoint "$n"); then
      echo "  (unknown ssh target)"
      continue
    fi
    # shellcheck disable=SC2086  # $t may include `-J bastion`; deliberate word-split.
    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."