#!/bin/sh
# Drives an etcd StatefulSet rolling upgrade with zero uncontrolled elections.
# Reads config from /config/config.yaml (mounted ConfigMap).
# Orchestration: move leader to -0, roll (N-1)..1, move leader to -(N-1), roll -0.
# Rolling is driven by decrementing spec.updateStrategy.rollingUpdate.partition
# (StatefulSet stays in RollingUpdate mode, locked at partition=replicas in
# steady state). The script never deletes pods; the sts controller does.
# Idempotent + crash-resumable: pods at sts.status.updateRevision are skipped,
# and partition is re-asserted to replicas on every exit path.
# Written to run under POSIX ash (alpine/k8s image uses busybox sh).

set -eu

CFG=/config/config.yaml
DIAL_TO="3s"

STS=$(yq         '.stsName'                   "$CFG")
NS=$(yq          '.namespace'                 "$CFG")
SVC=$(yq         '.serviceName'               "$CFG")
PORT=$(yq        '.clientPort // 2379'        "$CFG")
COOLDOWN=$(yq    '.cooldownSeconds // 30'     "$CFG")
READY_TO=$(yq    '.readyTimeoutSeconds // 300' "$CFG")
HEALTH_TO=$(yq   '.healthTimeoutSeconds // 60' "$CFG")

TLS_ON=$(yq '.tls.enabled // false' "$CFG")
SCHEME=http
ETCDCTL_TLS=""
if [ "$TLS_ON" = "true" ]; then
  SCHEME=https
  CA=$(yq   '.tls.caFile'   "$CFG")
  CERT=$(yq '.tls.certFile' "$CFG")
  KEY=$(yq  '.tls.keyFile'  "$CFG")
  ETCDCTL_TLS="--cacert=$CA --cert=$CERT --key=$KEY"
fi

REPLICAS=$(kubectl -n "$NS" get sts "$STS" -o jsonpath='{.spec.replicas}')
TARGET_REV=$(kubectl -n "$NS" get sts "$STS" -o jsonpath='{.status.updateRevision}')
if [ -z "$REPLICAS" ] || [ -z "$TARGET_REV" ]; then
  echo "FATAL: could not read replicas or updateRevision from sts $STS" >&2
  exit 1
fi

# Set spec.updateStrategy.rollingUpdate.partition. The sts controller deletes
# pods whose ordinal is >= partition and not yet at TARGET_REV; OrderedReady
# ensures it processes them one at a time. Idempotent: re-patching the same
# value is a no-op at the API server.
set_partition() {
  p=$1
  kubectl -n "$NS" patch sts "$STS" --type=merge \
    -p "{\"spec\":{\"updateStrategy\":{\"rollingUpdate\":{\"partition\":$p}}}}" >/dev/null
}

# Re-assert partition=$REPLICAS on every exit path (including failures after a
# partition decrement). Without this, a failure between `set_partition $i` and
# the final re-lock leaves the StatefulSet at a low partition, so the next
# unrelated template change could start an uncontrolled rolling update.
# Trap disarms itself first to avoid infinite recursion if set_partition fails.
relock_partition() {
  rc=$?
  trap - EXIT INT TERM
  set +e
  echo "[rollout] exit rc=$rc; re-asserting partition=$REPLICAS lock"
  set_partition "$REPLICAS" || echo "[rollout] WARN: re-lock partition failed; manual check required" >&2
  exit "$rc"
}
trap relock_partition EXIT INT TERM

# Build cluster-wide endpoints list once.
ENDPOINTS=""
i=0
while [ "$i" -lt "$REPLICAS" ]; do
  ep="${SCHEME}://${STS}-$i.${SVC}.${NS}:${PORT}"
  if [ -z "$ENDPOINTS" ]; then ENDPOINTS="$ep"; else ENDPOINTS="$ENDPOINTS,$ep"; fi
  i=$((i + 1))
done

ep_of() { echo "${SCHEME}://$1.${SVC}.${NS}:${PORT}"; }

# Run etcdctl inside <pod>, targeting that pod's own local endpoint (per-pod commands).
etcdctl_on() {
  pod=$1; shift
  kubectl -n "$NS" exec "$pod" -c etcd -- \
    sh -c "etcdctl --endpoints=$(ep_of "$pod") --dial-timeout=$DIAL_TO $ETCDCTL_TLS $*"
}

# Run etcdctl inside <pod>, targeting all cluster endpoints (cluster-view queries).
etcdctl_cluster_on() {
  pod=$1; shift
  kubectl -n "$NS" exec "$pod" -c etcd -- \
    sh -c "etcdctl --endpoints=$ENDPOINTS --dial-timeout=$DIAL_TO $ETCDCTL_TLS $*"
}

# Find a pod we can exec into. Tries -0, -1, ..., -(N-1). Prints name or fails.
pick_live_pod() {
  j=0
  while [ "$j" -lt "$REPLICAS" ]; do
    p="${STS}-$j"
    if kubectl -n "$NS" exec "$p" -c etcd -- true >/dev/null 2>&1; then
      echo "$p"
      return 0
    fi
    j=$((j + 1))
  done
  return 1
}

# Return the pod name of the current raft leader, or empty string if undetermined.
# endpoint status -w simple columns: endpoint,memberID,version,dbSize,isLeader,isLearner,...
# Unreachable endpoints produce stderr warnings + non-zero exit, but reachable rows
# still appear on stdout -- so we swallow the exit code and parse stdout.
get_leader_pod() {
  probe=$1
  out=$(etcdctl_cluster_on "$probe" endpoint status -w simple 2>/dev/null || true)
  echo "$out" | awk -F', ' '$5=="true" { print $1; exit }' \
    | sed -E 's|^https?://([^.]+)\..*$|\1|'
}

# Hex member ID for a named pod, or empty on failure.
# member list -w simple columns: ID,status,name,peerURLs,clientURLs,isLearner
member_id_of() {
  probe=$1; target=$2
  kubectl -n "$NS" exec "$probe" -c etcd -- \
    sh -c "etcdctl --dial-timeout=$DIAL_TO $ETCDCTL_TLS member list -w simple" 2>/dev/null \
    | awk -F', ' -v n="$target" '$3==n { print $1 }'
}

# Move raft leader to $1. No-op if already there.
# Retries a few times to cover the short window after target restart where log is still catching up.
move_leader_to() {
  target=$1
  attempt=1
  while [ "$attempt" -le 5 ]; do
    probe=$(pick_live_pod) || { echo "[rollout] ERROR: no live pod to probe" >&2; return 1; }
    leader=$(get_leader_pod "$probe")
    if [ -z "$leader" ]; then
      echo "[rollout] attempt $attempt: could not locate leader via $probe; retrying" >&2
      sleep 3; attempt=$((attempt + 1)); continue
    fi
    if [ "$leader" = "$target" ]; then
      echo "[rollout] leader already on $target"
      return 0
    fi
    id=$(member_id_of "$probe" "$target")
    if [ -z "$id" ]; then
      echo "[rollout] attempt $attempt: no member id for $target; retrying" >&2
      sleep 3; attempt=$((attempt + 1)); continue
    fi
    echo "[rollout] attempt $attempt: transfer leader $leader -> $target (id=$id)"
    if etcdctl_on "$leader" move-leader "$id"; then return 0; fi
    echo "[rollout] attempt $attempt: move-leader failed; retrying" >&2
    sleep 5; attempt=$((attempt + 1))
  done
  echo "[rollout] ERROR: gave up moving leader to $target" >&2
  return 1
}

pod_up_to_date() {
  pod=$1
  kubectl -n "$NS" get pod "$pod" >/dev/null 2>&1 || return 1
  rev=$(kubectl -n "$NS" get pod "$pod" \
    -o jsonpath='{.metadata.labels.controller-revision-hash}' 2>/dev/null || echo "")
  ready=$(kubectl -n "$NS" get pod "$pod" \
    -o jsonpath='{.status.conditions[?(@.type=="Ready")].status}' 2>/dev/null || echo "")
  [ "$rev" = "$TARGET_REV" ] && [ "$ready" = "True" ]
}

# Wait for <pod> to reach target revision + Ready, then pass etcdctl endpoint health.
# Timeouts are inline (not via `timeout` + `sh -c`) because POSIX sh does not export functions.
wait_pod_healthy() {
  pod=$1
  deadline=$(( $(date +%s) + READY_TO ))
  until pod_up_to_date "$pod"; do
    if [ "$(date +%s)" -ge "$deadline" ]; then
      echo "[rollout] ERROR: timed out waiting for $pod to reach $TARGET_REV + Ready" >&2
      return 1
    fi
    sleep 2
  done
  deadline=$(( $(date +%s) + HEALTH_TO ))
  until etcdctl_on "$pod" endpoint health >/dev/null 2>&1; do
    if [ "$(date +%s)" -ge "$deadline" ]; then
      echo "[rollout] ERROR: timed out waiting for $pod etcd health" >&2
      return 1
    fi
    sleep 2
  done
}

echo "[rollout] sts=$STS ns=$NS replicas=$REPLICAS target_rev=$TARGET_REV tls=$TLS_ON"

# ---- fast path: everything already at target revision ----
all_done=true
i=0
while [ "$i" -lt "$REPLICAS" ]; do
  if ! pod_up_to_date "${STS}-$i"; then all_done=false; break; fi
  i=$((i + 1))
done
if $all_done; then
  echo "[rollout] all pods already at $TARGET_REV; re-asserting partition=$REPLICAS lock"
  set_partition "$REPLICAS"
  exit 0
fi

# ---- wait until we can exec into a live pod (greenfield / post-restart) ----
# dependsOnSts only guarantees the StatefulSet object exists, not that any
# pod is Running and ready to exec. On a brand-new cluster (slow scheduling,
# slow image pull) pick_live_pod can fail for minutes; poll inside READY_TO
# instead of bailing on the first miss.
echo "[rollout] waiting for a live pod we can exec into"
probe=""
deadline=$(( $(date +%s) + READY_TO ))
while [ -z "$probe" ]; do
  if probe=$(pick_live_pod); then break; fi
  if [ "$(date +%s)" -ge "$deadline" ]; then
    echo "[rollout] FATAL: no live pod after ${READY_TO}s" >&2
    exit 1
  fi
  sleep 2
done

# ---- wait until cluster has a leader (greenfield / post-restart safety) ----
echo "[rollout] waiting for cluster to elect a leader (probe=$probe)"
deadline=$(( $(date +%s) + READY_TO ))
while [ -z "$(get_leader_pod "$probe")" ]; do
  if [ "$(date +%s)" -ge "$deadline" ]; then
    echo "[rollout] FATAL: cluster has no leader after ${READY_TO}s" >&2
    exit 1
  fi
  sleep 2
  # Refresh probe in case the original pod got replaced mid-wait.
  probe=$(pick_live_pod) || { echo "[rollout] FATAL: no live pod during leader wait" >&2; exit 1; }
done

# ---- 1) move leader to -0 ----
move_leader_to "${STS}-0"

# ---- 2) roll (N-1) .. 1 ----
i=$((REPLICAS - 1))
while [ "$i" -ge 1 ]; do
  pod="${STS}-$i"
  if pod_up_to_date "$pod"; then
    echo "[rollout] skip $pod (already $TARGET_REV)"
    i=$((i - 1))
    continue
  fi
  echo "[rollout] partition=$i (rolls $pod)"
  set_partition "$i"
  wait_pod_healthy "$pod"
  echo "[rollout] $pod healthy; cooldown ${COOLDOWN}s"
  sleep "$COOLDOWN"
  i=$((i - 1))
done

# ---- 3) move leader to -(N-1), then roll -0 ----
move_leader_to "${STS}-$((REPLICAS - 1))"

pod="${STS}-0"
if ! pod_up_to_date "$pod"; then
  echo "[rollout] partition=0 (rolls $pod)"
  set_partition 0
  wait_pod_healthy "$pod"
fi

# ---- 4) re-lock: bump partition back to replicas ----
echo "[rollout] re-asserting partition=$REPLICAS lock"
set_partition "$REPLICAS"

echo "[rollout] done"