#!/usr/bin/env bash
#
# ucp-etcd-v2store-mitigation.sh
#
# Interactive, one-manager-at-a-time mitigation for the MKE 3 -> MKE 4 (4.2.0)
# migration failure at [Migrate MKE 3 Nodes], where mkectl's v2store-clean check
# rejects the bundled etcdutl's clean output ("No custom content found in both
# v2store and WAL records.") and loops to timeout + rollback.
#
# Mitigation: on each manager, replace the running ucp-kv (etcd 3.5.32 in MKE
# 3.9.5) with the 3.9.2 image (etcd 3.5.21) by retagging over the configured
# tag and letting ucp-agent recreate the container. Recreating on 3.5.21
# rewrites the on-disk etcd state so the migration's v2store check passes.
#
# Run SEPARATELY on each manager, one at a time; only move to the next manager
# after this one is healthy and you have confirmed etcd quorum.
#
# Usage:
#   ./ucp-etcd-v2store-mitigation.sh [apply]     # apply on this manager (gated)
#   ./ucp-etcd-v2store-mitigation.sh rollback    # restore the original image
#   ./ucp-etcd-v2store-mitigation.sh cleanup     # only if you rolled back AND won't retry (see below)
#
#   --assume-yes | -y     auto-answer all gates (for automated end-to-end runs);
#                         DEFAULT is fully interactive/gated.
#
# Golden images are discovered, not configured: the fix and rollback images are
# taken from the SAME repo as the running ucp-kv (from `docker inspect ucp-kv`),
# i.e. <repo>:3.9.2 and <repo>:3.9.5. AIRGAP prerequisite: stage ucp-etcd:3.9.2
# into that same repo (where ucp-etcd:3.9.5 already lives) before running -- 3.9.2
# is not in the 3.9.5 bundle. Nothing is ever pushed to the registry by this
# script; the registry copy of :3.9.5 is only re-pulled as a rollback fallback.
#
# When is cleanup needed?
#   ONLY when a rollback was performed AND you do not intend to retry the upgrade
#   right away. In that case cleanup removes the rollback tag + state file so the
#   node is left in a clean original state.
#
#   It is NOT needed after a successful upgrade with no rollback: a successful
#   3->4 migration replaces etcd -- the node then runs k0s etcd under containerd
#   and the whole MKE3 docker layer (ucp-kv, ucp-agent) is gone, so the tags this
#   script created (the retagged <repo>:SOURCE_TAG and the timestamped rollback
#   tag) and the state file are inert leftovers that nothing references and that
#   cannot affect MKE4. There is nothing to undo.
#
# Retrying after a rollback:
#   Do NOT run cleanup; re-run 'apply' in full instead. apply re-does every step
#   from scratch, INCLUDING re-pointing <repo>:SOURCE_TAG at the fix image --
#   never assume the tag still points where a previous run left it, because it
#   may have been reversed outside this script. apply also refuses to run if the
#   image currently behind ucp-kv already reports the fix etcd version (so a tag
#   left pointing at the fix image is never mistaken for the original); restore
#   the true original first in that case.
#
set -Eeuo pipefail
umask 077   # state file holds image/rollback metadata; keep it owner-only

CONTAINER="ucp-kv"
SOURCE_TAG="3.9.5"          # the configured (affected) tag on this cluster
FIX_TAG="3.9.2"             # the image carrying etcd/etcdutl 3.5.21
EXPECTED_VERSION="3.5.21"   # etcdutl version the FIX_TAG image must report
DATA_DIR="/data/datav3"     # etcd v3 data dir inside ucp-kv
WAIT_SECONDS=300
STATE_FILE="/tmp/ucp-etcd-workaround.state"

# ------------------------------ arg parsing --------------------------------- #
ASSUME_YES=false
MODE=""
while [[ $# -gt 0 ]]; do
  case "$1" in
    --assume-yes|-y) ASSUME_YES=true ;;
    apply|rollback|cleanup)
      [[ -z "$MODE" ]] || { echo "ERROR: multiple modes given" >&2; exit 1; }
      MODE="$1" ;;
    *) echo "ERROR: unknown argument: $1" >&2; exit 1 ;;
  esac
  shift
done
MODE="${MODE:-apply}"

die() { echo "ERROR: $*" >&2; exit 1; }

confirm() {
  local message="$1"
  echo
  echo "STOP AND VERIFY:"
  echo "$message"
  if [[ "$ASSUME_YES" == true ]]; then
    echo "[--assume-yes] auto-continuing"
    return 0
  fi
  read -r -p "Type YES to continue: " answer
  [[ "$answer" == "YES" ]] || die "Stopped by operator"
}

# stdin need not be a tty in --assume-yes mode (automation); otherwise require it.
if [[ "$ASSUME_YES" != true && ! -t 0 ]]; then
  die "This script must be run interactively (or pass --assume-yes for automation)"
fi

if [[ "$EUID" -eq 0 ]]; then
  docker_cmd() { docker "$@"; }
else
  docker_cmd() { sudo docker "$@"; }
fi

container_health() {
  docker_cmd inspect --type container "$CONTAINER" \
    --format '{{if .State.Health}}{{.State.Health.Status}}{{else}}none{{end}}'
}

show_status() {
  docker_cmd inspect --type container "$CONTAINER" \
    --format 'status={{.State.Status}} health={{if .State.Health}}{{.State.Health.Status}}{{else}}not-configured{{end}} image={{.Image}}'
}

mke4_present() {
  command -v k0s >/dev/null 2>&1 && return 0
  systemctl is-active --quiet k0scontroller 2>/dev/null && return 0
  return 1
}

# Wait until the container is running AND healthy (or running with no healthcheck
# configured). ucp-agent recreates ucp-kv within seconds but passes through
# "health: starting" (~35s observed) before "healthy" -- waiting only for
# existence and then asserting health races and fails spuriously.
wait_for_ready() {
  local deadline=$((SECONDS + WAIT_SECONDS)) status health
  while (( SECONDS < deadline )); do
    if docker_cmd inspect --type container "$CONTAINER" >/dev/null 2>&1; then
      status="$(docker_cmd inspect --type container "$CONTAINER" --format '{{.State.Status}}')"
      health="$(container_health)"
      if [[ "$status" == "running" && ( "$health" == "healthy" || "$health" == "none" ) ]]; then
        return 0
      fi
    fi
    sleep 5
  done
  return 1
}

# Wait for readiness, verify the recreated container runs the expected image ID,
# and -- when no Docker healthcheck exists -- require an explicit operator
# confirmation of etcd/member health (existence/running is not readiness).
verify_recreated() {
  local expected_image_id="$1" running_id health
  echo "Waiting for ucp-agent to recreate $CONTAINER..."
  wait_for_ready || die "$CONTAINER was not running/healthy within ${WAIT_SECONDS}s"
  echo; echo "Recreated container:"; show_status

  running_id="$(docker_cmd inspect --type container "$CONTAINER" --format '{{.Image}}')"
  [[ "$running_id" == "$expected_image_id" ]] \
    || die "recreated $CONTAINER is not using the expected image ($expected_image_id)"
  [[ "$(docker_cmd inspect --type container "$CONTAINER" --format '{{.State.Status}}')" == "running" ]] \
    || die "$CONTAINER is not running"

  health="$(container_health)"
  if [[ "$health" == "none" ]]; then
    confirm "No Docker healthcheck is configured on $CONTAINER. Confirm OUT-OF-BAND
that etcd on this manager is up, has rejoined membership, and quorum is healthy
before continuing."
  fi
  docker_cmd exec "$CONTAINER" etcdutl version | head -1
}

# ------------------------------- rollback ----------------------------------- #
if [[ "$MODE" == "rollback" ]]; then
  [[ -f "$STATE_FILE" ]] || die "No state file at $STATE_FILE; nothing to roll back"
  # shellcheck disable=SC1090
  source "$STATE_FILE"
  : "${IMAGE_REF:?missing in state file}" "${ORIGINAL_TAG:?missing in state file}" "${ORIGINAL_IMAGE_ID:?missing in state file}"
  echo "Original image tag : $ORIGINAL_TAG"
  echo "Configured ref     : $IMAGE_REF"
  expected_id="$ORIGINAL_IMAGE_ID"
  if docker_cmd inspect --type image "$ORIGINAL_TAG" >/dev/null 2>&1; then
    # Preferred: the exact original image is still on this node.
    confirm "This restores the ORIGINAL image on this manager (points $IMAGE_REF back
at the locally-preserved $ORIGINAL_TAG and recreates $CONTAINER). Do this only if the
upgrade is delayed or has failed, and process one manager at a time."
    docker_cmd tag "$ORIGINAL_TAG" "$IMAGE_REF"
  else
    # Fallback: the preserved original was pruned. Re-pull $SOURCE_TAG from the
    # SAME repo as the cluster's configured ref (the registry copy is unmodified;
    # the script never pushed to it).
    golden_src="${IMAGE_REF%:*}:${SOURCE_TAG}"
    confirm "The locally-preserved original ($ORIGINAL_TAG) is GONE. Rollback will
re-pull $golden_src and point $IMAGE_REF at it, then recreate $CONTAINER.
Confirm $golden_src is the correct unmodified $SOURCE_TAG image for this cluster."
    docker_cmd pull "$golden_src" \
      || die "Cannot pull $golden_src for rollback. Ensure the $SOURCE_TAG image is
present in ${IMAGE_REF%:*}."
    docker_cmd run --rm --entrypoint etcdutl "$golden_src" version 2>&1 | head -1
    docker_cmd tag "$golden_src" "$IMAGE_REF"
    expected_id="$(docker_cmd inspect --type image "$golden_src" --format '{{.Id}}')"
  fi
  docker_cmd rm -f "$CONTAINER"
  verify_recreated "$expected_id"
  echo "Rollback complete. Verify etcd version and cluster/quorum health before"
  echo "restoring any other manager."
  exit 0
fi

# ------------------------------- cleanup ------------------------------------ #
# Needed ONLY when a rollback was performed AND you are not retrying the upgrade
# right away -- then this removes the rollback tag + state file to leave the node
# clean. NOT needed after a successful upgrade (etcd has moved to k0s/containerd
# and the MKE3 docker layer is gone, so the leftover tags/state file are inert).
# If you ARE retrying after a rollback, do NOT clean up -- re-run 'apply'. Run
# cleanup per-manager.
if [[ "$MODE" == "cleanup" ]]; then
  [[ -f "$STATE_FILE" ]] || die "No state file at $STATE_FILE; nothing to clean up"
  # shellcheck disable=SC1090
  source "$STATE_FILE"
  : "${IMAGE_REF:?missing in state file}" "${ORIGINAL_TAG:?missing in state file}"

  if docker_cmd inspect --type container "$CONTAINER" >/dev/null 2>&1; then
    running_id="$(docker_cmd inspect --type container "$CONTAINER" --format '{{.Image}}')"
    fix_id="$(docker_cmd inspect --type image "$IMAGE_REF" --format '{{.Id}}' 2>/dev/null || true)"
    if [[ -n "$fix_id" && "$running_id" == "$fix_id" ]]; then
      die "$CONTAINER is still running the temporary image. Do NOT clean up until the MKE4 upgrade has fully replaced etcd on this manager."
    fi
    echo "Note: $CONTAINER still exists but is not running the temporary image."
  else
    echo "$CONTAINER is gone (expected after a successful MKE4 upgrade)."
  fi

  if mke4_present; then
    echo "MKE4 detected on this node (k0s present)."
  else
    confirm "Could NOT detect k0s/MKE4 on this node. Confirm the MKE4 upgrade
completed successfully on THIS manager before removing the rollback tag."
  fi

  confirm "Confirm the MKE4 upgrade completed successfully on the ENTIRE cluster.
This removes the rollback tag $ORIGINAL_TAG and the state file $STATE_FILE
(tag removal only; the image is not force-deleted). Run cleanup on EVERY manager."

  # Tag removal only: 'docker rmi <tag>' refuses if a running container still
  # references the image, so this can never pull an in-use image out from under a
  # container. No -f.
  docker_cmd rmi "$ORIGINAL_TAG" 2>/dev/null || echo "rollback tag already absent or still referenced (left in place)"
  rm -f "$STATE_FILE"
  echo "Cleanup complete on this manager."
  exit 0
fi

# -------------------------------- apply ------------------------------------- #
confirm "Before continuing, confirm that:

- this is the intended manager;
- all managers currently have healthy etcd membership and quorum;
- a recent etcd backup exists;
- no other manager is being modified;
- the MKE4 upgrade can begin promptly after all managers pass;
- the $FIX_TAG image is from the approved registry."

echo "Reading the currently running image..."
IMAGE_REF="$(docker_cmd inspect --type container "$CONTAINER" --format '{{.Config.Image}}')"

case "$IMAGE_REF" in
  *@*) die "$CONTAINER uses an image digest. Stop and confirm ucp-agent behavior before continuing." ;;
esac

IMAGE_REPO="${IMAGE_REF%:*}"                       # the cluster's own repo (retag target)
EXPECTED_CONFIGURED_REF="${IMAGE_REPO}:${SOURCE_TAG}"
GOLDEN_FIX_REF="${IMAGE_REPO}:${FIX_TAG}"          # golden fix image, SAME repo as ucp-kv

[[ "$IMAGE_REF" == "$EXPECTED_CONFIGURED_REF" ]] \
  || die "Expected $EXPECTED_CONFIGURED_REF, found $IMAGE_REF"

ORIGINAL_IMAGE_ID="$(docker_cmd inspect --type container "$CONTAINER" --format '{{.Image}}')"
ORIGINAL_TAG="${IMAGE_REPO}:${SOURCE_TAG}-original-$(date +%Y%m%d%H%M%S)"

echo "Configured image  : $IMAGE_REF"
echo "Golden fix image  : $GOLDEN_FIX_REF"
echo "Original image ID : $ORIGINAL_IMAGE_ID"
echo "Rollback tag      : $ORIGINAL_TAG"

# Retry safety: the $SOURCE_TAG *ref* being present is not proof it still points
# at the true original image -- on a retry, someone may have swapped it outside
# this script. Verify the image currently behind ucp-kv reports the AFFECTED etcd
# version, so we never preserve the fix image as if it were the original.
echo "Confirming the current $CONTAINER image is the affected (pre-fix) etcd..."
CURRENT_ETCD_VER="$(docker_cmd run --rm --entrypoint etcdutl "$ORIGINAL_IMAGE_ID" version 2>&1 | head -1 || true)"
echo "Current $CONTAINER etcd: $CURRENT_ETCD_VER"
if grep -Eq "${EXPECTED_VERSION//./\\.}" <<<"$CURRENT_ETCD_VER"; then
  die "The current $CONTAINER image already reports etcd $EXPECTED_VERSION (the fix image).
$IMAGE_REF was swapped outside a clean state. Restore the true original image to
$IMAGE_REF first (e.g. '$0 rollback' if a state file exists, or retag the real
$SOURCE_TAG image), then re-run apply."
fi

# Pull + verify the golden fix image FIRST (fail-fast): if it is not available
# (e.g. airgap registry not staged with $FIX_TAG), abort here with the cluster
# completely unchanged -- nothing preserved, retagged, or removed yet.
echo "Pulling golden fix image $GOLDEN_FIX_REF ..."
docker_cmd pull "$GOLDEN_FIX_REF" \
  || die "Cannot pull $GOLDEN_FIX_REF. In airgap, stage the $FIX_TAG ucp-etcd image
into the same repo as the running ucp-kv (${IMAGE_REPO}). No changes made."
VERSION_OUTPUT="$(docker_cmd run --rm --entrypoint etcdutl "$GOLDEN_FIX_REF" version 2>&1)"
printf '%s\n' "$VERSION_OUTPUT"
grep -Eq "${EXPECTED_VERSION//./\\.}" <<<"$VERSION_OUTPUT" \
  || die "$GOLDEN_FIX_REF does not report etcdutl $EXPECTED_VERSION"
OLD_IMAGE_ID="$(docker_cmd inspect --type image "$GOLDEN_FIX_REF" --format '{{.Id}}')"
echo "golden fix image ID : $OLD_IMAGE_ID"
docker_cmd inspect --type image "$GOLDEN_FIX_REF" --format 'repo={{.RepoTags}} digests={{join .RepoDigests " "}}'

confirm "The golden fix image ($GOLDEN_FIX_REF) reports etcdutl $EXPECTED_VERSION.

Important: this image temporarily replaces the etcd SERVER as well as etcdutl.
The $EXPECTED_VERSION v2store checker may not inspect v2 WAL data as completely as
3.5.32 does -- confirm that is acceptable for this cluster's MKE store.
The current $SOURCE_TAG image is preserved locally as $ORIGINAL_TAG for rollback."

# Now that the fix image is confirmed available, preserve the original and record
# state. (Rollback re-derives the repo from IMAGE_REF to re-pull $SOURCE_TAG if the
# locally-preserved original is ever pruned.)
docker_cmd tag "$ORIGINAL_IMAGE_ID" "$ORIGINAL_TAG"
cat > "$STATE_FILE" <<EOF
IMAGE_REF=$IMAGE_REF
ORIGINAL_TAG=$ORIGINAL_TAG
ORIGINAL_IMAGE_ID=$ORIGINAL_IMAGE_ID
EOF

# Point the cluster's own $SOURCE_TAG at the golden fix image so ucp-agent uses it.
docker_cmd tag "$GOLDEN_FIX_REF" "$IMAGE_REF"

# GATE: before the destructive operation.
confirm "The local $IMAGE_REF tag now points to the golden $FIX_TAG image. The next
step REMOVES $CONTAINER so ucp-agent recreates it on etcd $EXPECTED_VERSION. Confirm
you are prepared for a brief etcd member interruption on THIS manager only."

docker_cmd rm -f "$CONTAINER"

verify_recreated "$OLD_IMAGE_ID"

# GATE: after container recreation.
confirm "Verify OUTSIDE this script that:

- this manager has rejoined the etcd cluster;
- etcd member and quorum health are good;
- the replacement container is stable;
- no other manager is currently being changed.

The running image ID was verified as $OLD_IMAGE_ID."

# Informational only; the result is not a pass condition. Both outcomes are
# normal: "detected custom content" before the migration has run (it deletes the
# v2 keyspace itself, then re-checks), "clean" on a cluster that already reached
# that step. What this proves is that the downgraded etcdutl works on the store.
set +e
CHECK_OUTPUT="$(docker_cmd exec --user nobody "$CONTAINER" \
  etcdutl check v2store --data-dir="$DATA_DIR" 2>&1)"
set -e
echo
echo "v2store check with etcdutl $EXPECTED_VERSION (informational):"
printf '%s\n' "$CHECK_OUTPUT"

# GATE: after the v2store check.
confirm "The v2store check above ran with etcdutl $EXPECTED_VERSION.

Because $EXPECTED_VERSION has older v2store-check behavior (it may not inspect v2
WAL data as completely as 3.5.32), confirm that relying on its result during the
migration is acceptable for this cluster's MKE store before proceeding."

# GATE: before moving to another manager.
confirm "This manager is complete. Confirm again that:

- this manager is healthy and etcd quorum is healthy;
- the NEXT manager is the only manager that will be modified next;
- the MKE4 upgrade will not be left waiting after all managers are processed.

The rollback state is saved in $STATE_FILE."

echo
echo "Manager complete."
echo "Run this script on the next manager only after the health checks above pass."
echo "If the upgrade is delayed or fails, run: $0 rollback"
echo "To retry after a rollback: re-run '$0' (apply) in full -- do NOT clean up."
echo "A successful upgrade needs no cleanup (etcd moves to k0s; leftover tags/state"
echo "are inert). Run '$0 cleanup' only if you rolled back and won't retry."