#!/usr/bin/env bash
# live swap (live swap): put a kernel image on a cluster that already exists and restart its pods.
#
# Patches the correct layer for the chosen lane, waits for the rollout, and prints which event will
# revert the change. Unlike the catalog path this writes a FULL image reference, so any registry the
# dataplane can pull from is allowed — including a registry the kernel team owns.
#
# See skills/get-started/references/change-kernel-images.md for the per-lane field map and
# the Starter traps (six TiKV partition CRs, pooled TiDB image lives in a ConfigMap).
set -Eeuo pipefail

usage() {
  cat >&2 <<'USAGE'
usage: patch-live-kernel-image.sh --lane <premium|dedicated|starter> --kind <kind> --name <name> \
                                  --namespace <ns> --image <full-image-ref> [options]

required:
  --lane premium|dedicated|starter
  --kind <kind>        premium/dedicated group CR: tidbgroup|tikvgroup|pdgroup|tiflashgroup|
                                                   ticdcgroup|tiproxygroup|tsogroup|schedulergroup
                       dedicated control plane   : configrevision|clusterconfig
                       starter                   : tidbcluster | statefulset | deployment | configmap
  --name <name>        object name
  --namespace <ns>     object namespace. Starter TidbCluster CRs live in tidb-serverless;
                       the pooled serverless-pool-config ConfigMap lives in the lane
                       namespace (tidb-pool, vip-tidb-pool, ...), not tidb-serverless.
                       Ignored for configrevision/clusterconfig, which are cluster-scoped.
  --image <ref>        FULL image reference, registry/repo:tag

options for specific kinds:
  --component <name>   required for configrevision|clusterconfig (pd|tidb|tikv|tiflash|tiproxy|...)
                       and for tidbcluster (pd|tidb|tikv|tiflash)
  --container <name>   required for statefulset|deployment
  --configmap-key <k>  key to patch for kind=configmap (default tidb-image)

connection:
  --kubeconfig <path>  kubeconfig for the target cluster (dataplane for premium/dedicated)
  --context <name>     kube context

behaviour:
  --skip-registry-check   skip verifying the image exists
  --aws-profile <p>       profile for the ECR existence check
  --rollout-timeout <d>   default 10m
  --apply                 perform the patch. Without this the script only previews.
USAGE
  exit 64
}

lane="" kind="" name="" namespace="" image="" component="" container=""
configmap_key="tidb-image" kubeconfig="" context="" aws_profile=""
skip_registry_check=false rollout_timeout="10m" apply=false

while (($#)); do
  case "$1" in
    --lane) lane="${2:-}"; shift 2 ;;
    --kind) kind=$(printf '%s' "${2:-}" | tr '[:upper:]' '[:lower:]'); shift 2 ;;
    --name) name="${2:-}"; shift 2 ;;
    --namespace) namespace="${2:-}"; shift 2 ;;
    --image) image="${2:-}"; shift 2 ;;
    --component) component="${2:-}"; shift 2 ;;
    --container) container="${2:-}"; shift 2 ;;
    --configmap-key) configmap_key="${2:-}"; shift 2 ;;
    --kubeconfig) kubeconfig="${2:-}"; shift 2 ;;
    --context) context="${2:-}"; shift 2 ;;
    --aws-profile) aws_profile="${2:-}"; shift 2 ;;
    --skip-registry-check) skip_registry_check=true; shift ;;
    --rollout-timeout) rollout_timeout="${2:-}"; shift 2 ;;
    --apply) apply=true; shift ;;
    -h|--help) usage ;;
    *) echo "unknown argument: $1" >&2; usage ;;
  esac
done

[[ -n "$lane" && -n "$kind" && -n "$name" && -n "$namespace" && -n "$image" ]] || usage
case "$lane" in premium|dedicated|starter) ;; *) echo "invalid --lane: $lane" >&2; exit 64 ;; esac

kubectl_args=(-n "$namespace")
[[ -n "$kubeconfig" ]] && kubectl_args=(--kubeconfig "$kubeconfig" "${kubectl_args[@]}")
[[ -n "$context" ]] && kubectl_args+=(--context "$context")
kc() { kubectl "${kubectl_args[@]}" "$@"; }

# ---- reject the layers that get reverted immediately --------------------------------------------
case "$lane:$kind" in
  premium:statefulset|premium:deployment|premium:pod|dedicated:statefulset|dedicated:deployment|dedicated:pod)
    cat >&2 <<'REJECT'
refusing: on the Premium and Dedicated lanes tidb-operator owns the StatefulSet/Deployment image
and reverts a direct patch almost immediately. Patch the group CR instead, for example:

  --kind tikvgroup --name <group> --namespace <cluster-ns>
REJECT
    exit 64
    ;;
  starter:pod)
    echo "refusing: pooled Starter pods are replaced from a template; patch the ConfigMap instead" >&2
    echo "  --kind configmap --name serverless-pool-config --configmap-key tidb-image" >&2
    exit 64
    ;;
esac

# ---- build the patch ----------------------------------------------------------------------------
# Split a reference into repository plus tag or digest. Handle the digest form first,
# then split repo:tag on the last colon that follows the last slash so a registry port
# (myreg:5000/tikv) is not mistaken for a tag.
image_repo="$image" image_tag="" image_digest=""
if [[ "$image" == *"@sha256:"* ]]; then
  image_repo=${image%@sha256:*}
  image_digest="sha256:${image##*@sha256:}"
else
  image_last_segment=${image##*/}
  if [[ "$image_last_segment" == *:* ]]; then
    image_tag=${image_last_segment##*:}
    image_repo=${image%":$image_tag"}
  fi
fi
[[ -n "$image_tag" || -n "$image_digest" ]] || {
  echo "--image must include a tag or a @sha256: digest" >&2
  exit 64
}

# Kinds that store repository and tag in separate fields cannot express a digest reference.
require_repo_tag_form() {
  [[ -z "$image_digest" ]] || {
    echo "kind $1 stores repository and tag in separate fields and cannot take a digest reference." >&2
    echo "Use a tag, or patch the dataplane group CR instead, which takes a full reference." >&2
    exit 64
  }
}

patch="" patch_type="merge" jsonpath="" rollout_target=""

case "$kind" in
  tidbgroup|tikvgroup|pdgroup|tiflashgroup|ticdcgroup|tiproxygroup|tsogroup|schedulergroup)
    patch=$(printf '{"spec":{"template":{"spec":{"image":"%s"}}}}' "$image")
    jsonpath='{.spec.template.spec.image}'
    ;;
  configrevision|clusterconfig)
    [[ -n "$component" ]] || { echo "--component is required for $kind" >&2; exit 64; }
    require_repo_tag_form "$kind"
    patch=$(printf '{"spec":{"%s":{"baseImage":"%s","tag":"%s"}}}' "$component" "$image_repo" "$image_tag")
    jsonpath="{.spec.${component}.baseImage}:{.spec.${component}.tag}"
    ;;
  tidbcluster)
    [[ -n "$component" ]] || { echo "--component is required for tidbcluster" >&2; exit 64; }
    require_repo_tag_form "$kind"
    patch=$(printf '{"spec":{"%s":{"baseImage":"%s","version":"%s"}}}' "$component" "$image_repo" "$image_tag")
    jsonpath="{.spec.${component}.baseImage}:{.spec.${component}.version}"
    ;;
  statefulset|deployment)
    [[ -n "$container" ]] || { echo "--container is required for $kind" >&2; exit 64; }
    patch=$(printf '{"spec":{"template":{"spec":{"containers":[{"name":"%s","image":"%s"}]}}}}' \
            "$container" "$image")
    patch_type="strategic"
    jsonpath="{.spec.template.spec.containers[?(@.name==\"${container}\")].image}"
    rollout_target="$kind/$name"
    ;;
  configmap)
    patch=$(printf '{"data":{"%s":"%s"}}' "$configmap_key" "$image")
    # Bracket form so a key containing dots (config.toml) is not read as a nested path.
    jsonpath="{.data['${configmap_key}']}"
    ;;
  *)
    echo "unsupported --kind: $kind" >&2
    exit 64
    ;;
esac

# ---- preflight ----------------------------------------------------------------------------------
kc get "$kind" "$name" >/dev/null || { echo "object not found: $kind/$name in $namespace" >&2; exit 1; }
current=$(kc get "$kind" "$name" -o "jsonpath=$jsonpath" 2>/dev/null || true)

if [[ "$skip_registry_check" == true ]]; then
  echo "warning: skipping registry check; a bad reference will only surface as ImagePullBackOff" >&2
elif [[ "$image_repo" == *.dkr.ecr.*.amazonaws.com/* ]]; then
  registry_host=${image_repo%%/*}
  repo_path=${image_repo#*/}
  region=${registry_host#*.dkr.ecr.}
  region=${region%%.amazonaws.com}
  if command -v aws >/dev/null; then
    if [[ -n "$image_digest" ]]; then
      image_id="imageDigest=$image_digest"
    else
      image_id="imageTag=$image_tag"
    fi
    aws_args=(ecr describe-images --region "$region" --repository-name "$repo_path"
              --image-ids "$image_id")
    [[ -n "$aws_profile" ]] && aws_args+=(--profile "$aws_profile")
    if aws "${aws_args[@]}" >/dev/null 2>&1; then
      echo "registry check ok: $repo_path ${image_id#*=} exists in $region"
    else
      echo "image not found in registry: $repo_path ${image_id#*=} (region $region)" >&2
      echo "push it first, or pass --skip-registry-check" >&2
      exit 1
    fi
  else
    echo "warning: aws CLI not found; skipping registry check" >&2
  fi
else
  echo "note: non-ECR registry, skipping existence check."
  echo "      Ensure the dataplane can pull it (cross-account pull grant for the runtime account)."
fi

# ---- revert warning -----------------------------------------------------------------------------
revert_note() {
  case "$lane" in
    premium)
      cat <<'NOTE'
cluster-ng-server recomputes this image from the metadb catalog only when a workflow requests
OpsUpdateImage, which happens for the VERSION and ALL config-version scopes
(cluster-service-ng domain/server/set_component_config_version.go:101-106) — a version change or
a full reconcile. Scale, quota, label, audit-log and config-only changes do NOT revert it.
For a change that survives a version change too, use a catalog swap (scripts/set-kernel-image.sh).
NOTE
      ;;
    dedicated)
      cat <<'NOTE'
cluster-controller re-renders ConfigRevision from Config Center on the next config or version
workflow, replacing this value with the catalog value.
For a change that survives that, use a catalog swap (scripts/set-kernel-image.sh --lane dedicated).
NOTE
      ;;
    starter)
      cat <<'NOTE'
No cp9 controller rewrites this continuously and the manager's version updater is off, BUT
`cp9 up` and `cp9 starter activate` re-apply the rendered overlay with
`kubectl apply --server-side --force-conflicts`
(personas/dev-full/starter/apply-remote-kustomize-images.sh:515), which takes field ownership back and
reverts this patch. Stable until the next activation.
Reminder: changing TiKV requires all SIX TidbCluster partition objects in
serverless-tikv-cluster.yaml; the tikv block in serverless-cluster.yaml has replicas: 0.
NOTE
      ;;
  esac
}

cat <<PREVIEW
lane        : $lane
object      : $kind/$name (namespace $namespace)
field       : $jsonpath
current     : ${current:-<unset>}
target      : $image
patch type  : $patch_type
patch       : $patch
PREVIEW
echo
echo "--- what will revert this ---"
revert_note

if [[ "$apply" != true ]]; then
  echo
  echo "dry run: nothing was patched. Re-run with --apply to perform the patch."
  exit 0
fi

# ---- apply --------------------------------------------------------------------------------------
kc patch "$kind" "$name" --type "$patch_type" -p "$patch"

after=$(kc get "$kind" "$name" -o "jsonpath=$jsonpath" 2>/dev/null || true)
echo "read-back: ${after:-<unset>}"
case "$kind" in
  configrevision|clusterconfig|tidbcluster)
    # these store repository and tag or version in separate fields
    expected="${image_repo}:${image_tag}"
    ;;
  *)
    expected="$image"
    ;;
esac
if [[ "$after" != "$expected" ]]; then
  echo "read-back mismatch: expected '$expected', got '${after:-<unset>}'" >&2
  echo "the patch did not land as written; a webhook or another controller may have rewritten it" >&2
  exit 1
fi

case "$kind" in
  configmap)
    cat <<EOF

ConfigMap patched. Pooled pods are NOT restarted automatically — delete them so replacements pick
up the new image, then confirm:

  kubectl ${kubectl_args[*]} delete pod -l <pool-selector>
  kubectl ${kubectl_args[*]} get pods -o jsonpath='{range .items[*]}{.metadata.name}{"\t"}{.spec.containers[0].image}{"\n"}{end}'
EOF
    ;;
  statefulset|deployment)
    echo "waiting for rollout of $rollout_target (timeout $rollout_timeout)"
    kc rollout status "$rollout_target" --timeout "$rollout_timeout"
    ;;
  *)
    cat <<EOF

The operator owns the rollout for this CR; it rolls the pods itself. Watch it converge:

  kubectl ${kubectl_args[*]} get pods -w
  kubectl ${kubectl_args[*]} get pods -o jsonpath='{range .items[*]}{.metadata.name}{"\t"}{.spec.containers[0].image}{"\n"}{end}'

Do not delete pods by hand to force it.
EOF
    ;;
esac

cat <<'DONE'

Record in the PR or run notes: the object patched, the full image reference, the observed running
pod image, and the revert exposure printed above. An unrecorded live patch is indistinguishable
from drift to the next person who inspects this environment.
DONE