#!/usr/bin/env bash
#
# eks-pause-resume.sh
#
# Pause / resume an EKS freetier cluster's COMPUTE by scaling every nodegroup to
# 0 and back, WITHOUT touching workloads (no kubectl, no replica bookkeeping).
# Workload specs (STS/Deployment) + PVCs live in the EKS-managed control plane,
# so they survive; on resume pods reschedule and EBS PVs reattach by volumeId.
#
# WHY NO CA STOP: the cluster-autoscaler can only SetDesiredCapacity within
# [min,max]; its IAM has no UpdateAutoScalingGroup. min=max=0 is a hard lock it
# cannot undo, and killing its node makes its pod Pending. It self-neutralizes
# and auto-recovers on resume. (We do NOT touch the CA at all.)
#
# WHY WE SCALE THE ASG DIRECTLY (not `eks update-nodegroup-config`): the EKS
# managed-nodegroup API rejects maxSize=0 ("valid min value: 1"), so it cannot
# express the lock. We therefore set the UNDERLYING ASG's min/max/desired with
# `aws autoscaling update-auto-scaling-group`, which accepts 0/0/0. The CA reads
# the ASG's bounds directly, so max=0 still locks it. This diverges the EKS
# nodegroup's recorded scalingConfig (it keeps showing the old max) from the live
# ASG — cosmetic only; verified EKS does NOT revert the ASG, and the next pulumi
# apply / nodegroup update re-syncs it. Resume restores via the ASG too (because
# the EKS record is unchanged, an `eks update-nodegroup-config` resume could be a
# no-op and never reach the ASG).
#
# RESTORE SOURCE: a committed snapshot at pause-stacks/<cluster>.json, generated
# by `snapshot` from the Pulumi stack config (main HEAD) via
# hack/derive-nodegroup-bounds.mjs, with live values as fallback for nodegroups
# the config doesn't describe (serverless TiKV + test-infra tiw-*). pulumi itself
# does NOT restore nodegroup bounds (managed nodegroup scalingConfig is
# ignoreChanges + min/max set by a trigger-gated local.Command), so this snapshot
# is the restore source.
#
# PRECONDITIONS (NOT checked here — you must):
#   * TiKV is on EBS (instance-store i7ie/i8g fully drained) before pausing.
#   * pulumi is NOT applying this stack, and no GitOps reconciles these ASGs.
#
# COMMANDS:
#   list      read-only: show every nodegroup + min/max/desired + scale-in prot.
#   snapshot  read-only on cloud: derive bounds from --config, merge with live,
#             write pause-stacks/<cluster>.json (commit this before pausing).
#   pause     scale every nodegroup to min=max=desired=0 (clears scale-in
#             protection first). Requires the snapshot to already exist.
#   resume    restore min/max from the snapshot. desired = max(min, live so a
#             retry does not scale in). When min=0, seed default-burstable to
#             desired=1 so CA can schedule and refill the rest.
#
# SAFETY: dry-run is the DEFAULT; mutations run only with --apply. list/snapshot
# never mutate cloud state.
#
# USAGE:
#   hack/eks-pause-resume.sh --profile qa-sso --cluster dev-us-east-1-f02 list
#   hack/eks-pause-resume.sh --profile qa-sso --cluster dev-us-east-1-f02 \
#       --config pool/eks-freetier/Pulumi.dev-us-east-1-f02.yaml snapshot
#   hack/eks-pause-resume.sh --profile qa-sso --cluster dev-us-east-1-f02 pause          # dry-run
#   hack/eks-pause-resume.sh --profile qa-sso --cluster dev-us-east-1-f02 pause  --apply
#   hack/eks-pause-resume.sh --profile qa-sso --cluster dev-us-east-1-f02 resume         # dry-run
#   hack/eks-pause-resume.sh --profile qa-sso --cluster dev-us-east-1-f02 resume --apply
#
# FLAGS:
#   --profile <p>   Optional AWS CLI profile. Omit when credentials come from
#                   environment variables, such as GitHub Actions secrets.
#   --cluster <c>   EKS cluster name == stack name (REQUIRED).
#   --region <r>    AWS region (default us-east-1).
#   --config <f>    Pulumi.<stack>.yaml — required for `snapshot`.
#   --config-ref <r>
#                   Revision that supplied --config; stored in the snapshot.
#   --state <f>     snapshot path (default: <repo>/pause-stacks/<cluster>.json).
#   --apply         actually execute mutations (default: dry-run).
#
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)"

PROFILE="" CLUSTER="" REGION="us-east-1" CONFIG="" CONFIG_REF="" STATE="" APPLY=0 CMD=""

die() { echo "ERROR: $*" >&2; exit 1; }
usage() { sed -n '2,67p' "$0"; }

while [[ $# -gt 0 ]]; do
  case "$1" in
    --profile) PROFILE="${2:-}"; shift 2;;
    --cluster) CLUSTER="${2:-}"; shift 2;;
    --region)  REGION="${2:-}"; shift 2;;
    --config)  CONFIG="${2:-}"; shift 2;;
    --config-ref) CONFIG_REF="${2:-}"; shift 2;;
    --state)   STATE="${2:-}"; shift 2;;
    --apply)   APPLY=1; shift;;
    list|snapshot|pause|resume) CMD="$1"; shift;;
    -h|--help) usage; exit 0;;
    *) die "unknown arg: $1";;
  esac
done

[[ -n "$CLUSTER" ]] || die "--cluster required"
[[ -n "$CMD"     ]] || die "command required: list | snapshot | pause | resume"
[[ -n "$STATE"   ]] || STATE="$REPO_ROOT/pause-stacks/${CLUSTER}.json"
command -v jq >/dev/null || die "jq is required"

AWS=(aws --region "$REGION")
[[ -n "$PROFILE" ]] && AWS+=(--profile "$PROFILE")
CLUSTER_TAG="kubernetes.io/cluster/${CLUSTER}"

banner() {
  if [[ "$APPLY" == 1 ]]; then echo ">>> APPLY (will mutate)"; else echo ">>> DRY-RUN (no changes; --apply to execute)"; fi
  echo ">>> cmd=$CMD cluster=$CLUSTER region=$REGION profile=$PROFILE"
  echo ">>> state=$STATE"; echo
}

FAILED=0
# Run a mutating aws call: print always; execute only under --apply. On failure
# do NOT abort the whole batch (one transient AWS error must not leave a stack
# half-paused); count it and keep going. cmd_pause/cmd_resume exit non-zero if any.
run() {
  if [[ "$APPLY" == 1 ]]; then
    echo "+ $*" >&2
    if ! "$@"; then echo "  !! FAILED: $*" >&2; FAILED=$((FAILED + 1)); fi
  else
    echo "DRYRUN  $*"
  fi
}

# One compact JSON line per cluster ASG (read-only).
discover() {
  "${AWS[@]}" autoscaling describe-auto-scaling-groups --output json --query 'AutoScalingGroups[]' \
  | jq -c --arg tag "$CLUSTER_TAG" '
      .[] | select(any(.Tags[]?; .Key==$tag and .Value=="owned"))
      | { asg:.AutoScalingGroupName,
          ng:((.Tags[]?|select(.Key=="eks:nodegroup-name")|.Value)//null),
          name:((.Tags[]?|select(.Key=="eks:nodegroup-name")|.Value)//.AutoScalingGroupName),
          min:.MinSize, max:.MaxSize, desired:.DesiredCapacity,
          protectedInstances:[ .Instances[]?|select(.ProtectedFromScaleIn==true)|.InstanceId ] }'
}

cmd_list() {
  banner
  local rows; rows="$(discover)"
  [[ -n "$rows" ]] || die "no ASGs for ${CLUSTER_TAG}=owned (check cluster/region/profile)"
  { printf 'NODEGROUP\tMIN\tMAX\tDESIRED\t#PROTECTED\n'
    jq -r '[.name,(.min|tostring),(.max|tostring),(.desired|tostring),(.protectedInstances|length|tostring)]|@tsv' <<<"$rows"
  } | column -t -s $'\t'
  echo; echo "total nodegroups: $(jq -rs length <<<"$rows")  live nodes: $(jq -rs 'map(.desired)|add' <<<"$rows")"
}

cmd_snapshot() {
  banner
  [[ -n "$CONFIG" ]] || die "snapshot requires --config <Pulumi.<stack>.yaml>"
  [[ -f "$CONFIG" ]] || die "config not found: $CONFIG"
  command -v node >/dev/null || die "node is required for snapshot"

  echo ">>> deriving intended bounds from $CONFIG (${CONFIG_REF:-current checkout})"
  local cfgbounds; cfgbounds="$(node "$SCRIPT_DIR/derive-nodegroup-bounds.mjs" --config "$CONFIG")"
  local live; live="$(discover | jq -s '.')"
  [[ "$(jq 'length' <<<"$live")" -gt 0 ]] || die "no live ASGs discovered"

  local ref; ref="${CONFIG_REF:-$(git -C "$REPO_ROOT" rev-parse --short HEAD 2>/dev/null || echo unknown)}"
  local now; now="$(date -u +%FT%TZ)"

  # config takes precedence; live as fallback where config is silent.
  # If live is already locked (0/0) — e.g. re-snapshot after pause — keep the
  # previous snapshot's bounds so resume is not written as 0/0.
  local prev="{}"
  if [[ -f "$STATE" ]]; then
    prev="$(jq '.nodegroups // {}' "$STATE")"
  fi
  local merged
  merged="$(jq -n --argjson cfg "$cfgbounds" --argjson live "$live" --argjson prev "$prev" '
    [ $live[] | . as $l | ($cfg[$l.name]) as $c | ($prev[$l.name]) as $p
      | (
          if $c then {min:$c.min, max:$c.max, source:"config"}
          elif ($l.min == 0 and $l.max == 0 and $p and ($p.max // 0) > 0) then
            {min:$p.min, max:$p.max, source:"snapshot-preserve"}
          else {min:$l.min, max:$l.max, source:"live-fallback"}
          end
        ) as $b
      | { name:$l.name, asg:$l.asg, ng:$l.ng,
          min:$b.min, max:$b.max, source:$b.source,
          liveMin:$l.min, liveMax:$l.max }
    ] | sort_by(.name)')"

  mkdir -p "$(dirname "$STATE")"
  jq -n --arg cl "$CLUSTER" --arg rg "$REGION" --arg pf "$PROFILE" \
        --arg cfg "$CONFIG" --arg ref "$ref" --arg now "$now" --argjson ng "$merged" '
    { cluster:$cl, region:$rg, profile_hint:$pf,
      config: $cfg, config_ref:$ref, generated_at:$now,
      note: "min/max for resume; desired restored = min, CA refills. source=config or live-fallback (config silent: serverless TiKV + tiw test-infra).",
      nodegroups: ( $ng | map({ (.name): {min:.min, max:.max, source:.source, ng:.ng, asg:.asg} }) | add ) }' \
    > "$STATE"

  echo ">>> wrote $STATE"
  echo "    nodegroups: $(jq '.nodegroups|length' "$STATE")  (config: $(jq -r '[.nodegroups[]|select(.source=="config")]|length' "$STATE"), live-fallback: $(jq -r '[.nodegroups[]|select(.source=="live-fallback")]|length' "$STATE"), snapshot-preserve: $(jq -r '[.nodegroups[]|select(.source=="snapshot-preserve")]|length' "$STATE"))"
  echo
  echo "=== live-fallback (config silent) ==="
  jq -r '.[]|select(.source=="live-fallback")|"  \(.name)  min=\(.min) max=\(.max)"' <<<"$merged"
  echo "=== snapshot-preserve (live was 0/0) ==="
  jq -r '.[]|select(.source=="snapshot-preserve")|"  \(.name)  min=\(.min) max=\(.max)"' <<<"$merged"
  echo "=== drift normalized to config (live -> config) ==="
  jq -r '.[]|select(.source=="config" and (.min!=.liveMin or .max!=.liveMax))|"  \(.name)  live(\(.liveMin)/\(.liveMax)) -> config(\(.min)/\(.max))"' <<<"$merged"
  echo; echo ">>> review this file, commit it to pause-stacks/, THEN run pause."
}

cmd_pause() {
  banner
  [[ -f "$STATE" ]] || die "snapshot $STATE missing — run 'snapshot' and commit it before pausing"
  local rows; rows="$(discover)"
  [[ -n "$rows" ]] || die "no live ASGs discovered (check cluster/region/profile)"

  # Every live nodegroup must be covered by the snapshot, else resume (which
  # iterates the snapshot) can never bring it back. A snapshot taken before a
  # live change is stale and must be regenerated.
  local uncovered
  uncovered="$(comm -23 <(jq -r '.name' <<<"$rows" | sort -u) <(jq -r '.nodegroups|keys[]' "$STATE" | sort -u))"
  if [[ -n "$uncovered" ]]; then
    { echo "live nodegroups NOT in snapshot $STATE:"; echo "$uncovered" | sed 's/^/  /'; } >&2
    die "snapshot is stale — re-run 'snapshot' (and commit) so every live nodegroup can be restored"
  fi

  echo "Scaling every ASG to min=max=desired=0 (clearing scale-in protection first):"
  while IFS= read -r r; do
    [[ -n "$r" ]] || continue
    local asg; asg=$(jq -r '.asg' <<<"$r")
    echo "* $(jq -r '.name' <<<"$r")"
    mapfile -t prot < <(jq -r '.protectedInstances[]?' <<<"$r")
    if [[ "${#prot[@]}" -gt 0 ]]; then
      run "${AWS[@]}" autoscaling set-instance-protection --auto-scaling-group-name "$asg" \
          --no-protected-from-scale-in --instance-ids "${prot[@]}"
    fi
    # ASG-direct: the EKS managed-nodegroup API rejects maxSize=0; the ASG accepts it.
    run "${AWS[@]}" autoscaling update-auto-scaling-group --auto-scaling-group-name "$asg" \
        --min-size 0 --max-size 0 --desired-capacity 0
  done <<<"$rows"
  echo
  if [[ "$FAILED" -gt 0 ]]; then echo ">>> $FAILED command(s) FAILED — re-run pause to retry the rest."; exit 1; fi

  # Postcondition (apply only): re-describe and fail if anything is not fully
  # locked — e.g. an instance that launched and got scale-in-protected between
  # discover and the scale-down would survive max=0. Re-running pause clears it.
  if [[ "$APPLY" == 1 ]]; then
    local notlocked
    notlocked="$(discover | jq -r 'select(.max>0 or (.protectedInstances|length>0)) | .name')"
    if [[ -n "$notlocked" ]]; then
      { echo ">>> NOT fully locked (max>0 or still-protected instances) — re-run pause:"; echo "$notlocked" | sed 's/^/  /'; } >&2
      exit 1
    fi
  fi
  echo "Done. CA dies with its node and cannot re-expand (ASG max=0)."
}

cmd_resume() {
  banner
  [[ -f "$STATE" ]] || die "snapshot $STATE not found — cannot restore without it"
  # Staging ledger: refuse to stomp a running cluster (desired=min would drain
  # nodegroups CA already refilled). Local snapshots without .lifecycle are ok.
  if jq -e '.lifecycle.state? | type == "string"' "$STATE" >/dev/null; then
    local st; st="$(jq -r '.lifecycle.state' "$STATE")"
    [[ "$st" == "paused" ]] || die "ledger state is '$st' (expected paused) — refusing to resume"
  fi
  # Capture discovery explicitly: with `discover | while`, a failed AWS/profile
  # call would silently yield an empty LIVE_ASG and skip every nodegroup while
  # exiting 0. Capturing into a var lets us fail loudly on empty discovery.
  local live; live="$(discover)"
  [[ -n "$live" ]] || die "no live ASGs discovered (check cluster/region/profile) — refusing to resume blindly"
  # Map each snapshot nodegroup to its CURRENT ASG (names can change if a
  # nodegroup was recreated); ones absent from live are skipped.
  declare -A LIVE_ASG LIVE_DESIRED
  while IFS=$'\t' read -r nm asg des; do
    LIVE_ASG["$nm"]="$asg"
    LIVE_DESIRED["$nm"]="$des"
  done < <(jq -r '[.name,.asg,(.desired|tostring)]|@tsv' <<<"$live")

  # CA has no toleration for exclusive taints; default-burstable is exclusive:false.
  is_ca_seed() { [[ "$1" == *-default-burstable-* ]]; }

  local zero_mins total_ng has_seed=0
  zero_mins="$(jq '[.nodegroups[] | select(.min == 0)] | length' "$STATE")"
  total_ng="$(jq '.nodegroups | length' "$STATE")"
  while IFS=$'\t' read -r name minv maxv; do
    if is_ca_seed "$name" && [[ -n "${LIVE_ASG[$name]:-}" ]] && (( maxv > 0 )); then
      has_seed=1
    fi
  done < <(jq -r '.nodegroups | to_entries[] | [.key, (.value.min|tostring), (.value.max|tostring)] | @tsv' "$STATE")
  if [[ "$zero_mins" == "$total_ng" && "$has_seed" -eq 0 ]]; then
    die "all restore mins are 0 and no live default-burstable seed ASG — CA cannot scale from zero"
  fi

  echo "Restoring min/max from snapshot (desired=max(min,live); seed default-burstable if min=0):"
  local restored=0 missing=0 seeded=0
  while IFS=$'\t' read -r name minv maxv; do
    local asg="${LIVE_ASG[$name]:-}"
    if [[ -z "$asg" ]]; then
      echo "! $name  -> SKIP (not present live; nodegroup removed since snapshot)"; missing=$((missing + 1)); continue
    fi
    local live_des="${LIVE_DESIRED[$name]:-0}"
    local desired="$minv"
    (( live_des > desired )) && desired="$live_des"
    if is_ca_seed "$name" && (( maxv > 0 )) && (( desired < 1 )); then
      desired=1
      seeded=$((seeded + 1))
    fi
    echo "* $name  -> min=$minv max=$maxv desired=$desired"
    run "${AWS[@]}" autoscaling update-auto-scaling-group --auto-scaling-group-name "$asg" \
        --min-size "$minv" --max-size "$maxv" --desired-capacity "$desired"
    restored=$((restored + 1))
  done < <(jq -r '.nodegroups | to_entries[] | [.key, (.value.min|tostring), (.value.max|tostring)] | @tsv' "$STATE")
  [[ "$missing" -gt 0 ]] && echo ">>> $missing snapshot nodegroup(s) skipped (no longer live)."
  [[ "$restored" -eq 0 ]] && die "no snapshot nodegroup matched a live ASG — discovery/cluster mismatch, nothing restored"
  echo
  if [[ "$FAILED" -gt 0 ]]; then echo ">>> $FAILED command(s) FAILED — re-run resume to retry the rest."; exit 1; fi
  echo "Done. Seeded $seeded default-burstable ASG(s); CA reschedules and refills."
}

case "$CMD" in
  list)     cmd_list;;
  snapshot) cmd_snapshot;;
  pause)    cmd_pause;;
  resume)   cmd_resume;;
  *) die "unknown command $CMD";;
esac