#!/bin/bash
# update-node.sh — drain, apt-update, reboot-if-needed, wait, uncordon a single node.
#
# Usage:
#   update-node.sh <node-name>
#
# Runs from an operator machine (or from the control plane); needs kubectl and
# ssh access to root@<node-ssh-target>. Node → SSH target resolution is in
# `node_ssh_target` below; edit that mapping when you add nodes.
#
# NEVER touches k3s config, k3s services, containerd, or any manifest.
# Only fixes it will attempt: apt/dpkg recovery on the same node (see step 3).
#
# Environment overrides:
#   DRAIN_TIMEOUT_SECONDS (default 600)
#   REBOOT_MAX_WAIT_SECONDS (default 600)
#   POST_UNCORDON_WAIT_SECONDS (default 180)
#   CONTROL_PLANE_HOST (default 178.105.17.239)
#   SSH_OPTS (default "-o ConnectTimeout=10 -o StrictHostKeyChecking=accept-new")
#   ASSUME_YES=1 to skip interactive confirmations
set -euo pipefail

NODE="${1:-}"
if [ -z "$NODE" ]; then
  echo "usage: $0 <node-name>" >&2
  exit 2
fi

DRAIN_TIMEOUT_SECONDS="${DRAIN_TIMEOUT_SECONDS:-600}"
REBOOT_MAX_WAIT_SECONDS="${REBOOT_MAX_WAIT_SECONDS:-600}"
POST_UNCORDON_WAIT_SECONDS="${POST_UNCORDON_WAIT_SECONDS:-180}"
CONTROL_PLANE_HOST="${CONTROL_PLANE_HOST:-178.105.17.239}"
SSH_OPTS="${SSH_OPTS:--o ConnectTimeout=10 -o StrictHostKeyChecking=accept-new}"

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

# node -> ssh target. Private-IP workers are reached via ProxyJump through the CP.
node_ssh_target() {
  case "$1" in
    k3s-cp-1)            echo "root@178.105.17.239" ;;
    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-update-runner)   echo "root@167.233.79.65" ;;
    *)                   die "unknown node $1 — update node_ssh_target() in $0" ;;
  esac
}

SSH_TARGET=$(node_ssh_target "$NODE")

log "=== update-node.sh $NODE ==="
log "ssh target: $SSH_TARGET"

# --- 1. pre-check --------------------------------------------------------------
log "[1/7] pre-check"
kubectl get node "$NODE" >/dev/null || die "node $NODE not found in cluster"
READY=$(kubectl get node "$NODE" -o jsonpath='{.status.conditions[?(@.type=="Ready")].status}')
[ "$READY" = "True" ] || die "node $NODE is not Ready before we start (Ready=$READY)"

for cond in DiskPressure MemoryPressure PIDPressure; do
  v=$(kubectl get node "$NODE" -o jsonpath="{.status.conditions[?(@.type==\"$cond\")].status}")
  [ "$v" = "True" ] && die "node $NODE has $cond=True — refuse to update"
done

# --- 2. cordon + drain --------------------------------------------------------
log "[2/7] cordon + drain (timeout ${DRAIN_TIMEOUT_SECONDS}s)"
kubectl cordon "$NODE"

set +e
kubectl drain "$NODE" \
  --ignore-daemonsets \
  --delete-emptydir-data \
  --timeout="${DRAIN_TIMEOUT_SECONDS}s"
DRAIN_RC=$?
set -e

if [ $DRAIN_RC -ne 0 ]; then
  log "drain FAILED (rc=$DRAIN_RC). Never force. Uncordoning $NODE and marking SKIPPED."
  kubectl uncordon "$NODE" || true
  echo "SKIPPED_DRAIN_FAILED $NODE"
  exit 3
fi

# --- 3. apt update on the node ------------------------------------------------
log "[3/7] apt update/upgrade on $NODE"
REMOTE_APT=$(cat <<'REMOTE'
set -euo pipefail
export DEBIAN_FRONTEND=noninteractive
APT_OPTS='-y -o Dpkg::Options::=--force-confdef -o Dpkg::Options::=--force-confold'

# Recover from any half-finished dpkg state before touching apt.
if ! dpkg --audit | grep -qE .; then
  :
else
  echo "dpkg audit reported issues, running dpkg --configure -a"
  dpkg --configure -a || true
fi

apt-get update

# Try upgrade; on broken deps, one attempt at apt-get -f install then retry.
if ! apt-get $APT_OPTS upgrade; then
  echo "upgrade failed, attempting apt-get -f install"
  apt-get $APT_OPTS -f install
  apt-get $APT_OPTS upgrade
fi

apt-get $APT_OPTS dist-upgrade
apt-get $APT_OPTS autoremove --purge
apt-get clean

if [ -f /var/run/reboot-required ]; then
  echo "REBOOT_REQUIRED=yes"
  echo "REBOOT_REASON<<EOF"
  cat /var/run/reboot-required.pkgs 2>/dev/null || echo "(no package list)"
  echo "EOF"
else
  echo "REBOOT_REQUIRED=no"
fi
REMOTE
)

APT_OUT=$(ssh $SSH_OPTS $SSH_TARGET "bash -s" <<< "$REMOTE_APT")
echo "$APT_OUT" | sed 's/^/    /'

if echo "$APT_OUT" | grep -q '^REBOOT_REQUIRED=yes'; then
  REBOOT=1
else
  REBOOT=0
fi

# --- 4. reboot if required ----------------------------------------------------
if [ "$REBOOT" -eq 1 ]; then
  log "[4/7] reboot required — rebooting $NODE"
  ssh $SSH_OPTS $SSH_TARGET 'systemctl reboot' || true

  # Give SSH a moment to actually drop before we start polling.
  sleep 15

  deadline=$(( $(date +%s) + REBOOT_MAX_WAIT_SECONDS ))
  while [ $(date +%s) -lt $deadline ]; do
    if ssh $SSH_OPTS -o ConnectTimeout=5 $SSH_TARGET 'uptime' >/dev/null 2>&1; then
      log "  $NODE ssh is back"
      break
    fi
    sleep 10
  done

  if ! ssh $SSH_OPTS -o ConnectTimeout=5 $SSH_TARGET 'uptime' >/dev/null 2>&1; then
    die "node $NODE did not return within ${REBOOT_MAX_WAIT_SECONDS}s — escalate"
  fi
else
  log "[4/7] no reboot needed"
fi

# --- 5. wait for kubelet Ready ------------------------------------------------
log "[5/7] wait for kubelet Ready on $NODE"
deadline=$(( $(date +%s) + 300 ))
while [ $(date +%s) -lt $deadline ]; do
  READY=$(kubectl get node "$NODE" -o jsonpath='{.status.conditions[?(@.type=="Ready")].status}' 2>/dev/null || echo Unknown)
  [ "$READY" = "True" ] && break
  sleep 5
done
[ "$READY" = "True" ] || die "kubelet on $NODE never returned Ready — escalate (do NOT change k3s config)"
log "  Ready=True"

# --- 6. uncordon --------------------------------------------------------------
log "[6/7] uncordon $NODE"
kubectl uncordon "$NODE"

# --- 7. post-node settle ------------------------------------------------------
log "[7/7] post-node settle (${POST_UNCORDON_WAIT_SECONDS}s) + health check"
sleep "$POST_UNCORDON_WAIT_SECONDS"

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
if RETRY_ON_TRANSIENT=1 "$SCRIPT_DIR/cluster-health.sh"; then
  log "=== $NODE update: OK ==="
else
  die "cluster health failed after updating $NODE — STOP the cycle, do NOT continue"
fi