#!/bin/bash
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
# shellcheck disable=SC1091
. "$SCRIPT_DIR/common.sh"

usage() {
  cat >&2 <<'USAGE'
Usage:
  pause_cluster.sh [--plan|--apply]

Runs on the control/build host. Stops every cse-worker, the local TiUP cluster,
then EC2. S3, IAM, EBS, security groups, and TiUP metadata are preserved.
Defaults to plan.
USAGE
}

action="${1:---plan}"
case "$action" in
  --plan) apply=0 ;;
  --apply) apply=1 ;;
  --help) usage; exit 0 ;;
  *) usage; exit 1 ;;
esac

STATE_FILE="${STATE_FILE:-}"
require_env STATE_FILE
unset_vars \
  PROFILE EXPECTED_AWS_ACCOUNT_ID REGION CLUSTER_NAME ALL_INSTANCE_IDS \
  S3_BUCKET S3_PREFIX TIKV_IAM_ROLE RUN_LOG_DIR TIUP_CLUSTER_DEPLOYED \
  PAUSE_WORKERS_STOPPED PAUSE_TIUP_STOPPED CLUSTER_PAUSED PROVISION_TOKEN
load_state_file "$STATE_FILE"
require_env PROFILE EXPECTED_AWS_ACCOUNT_ID REGION CLUSTER_NAME ALL_INSTANCE_IDS S3_BUCKET S3_PREFIX TIKV_IAM_ROLE RUN_LOG_DIR TIUP_CLUSTER_DEPLOYED PROVISION_TOKEN
[ "$TIUP_CLUSTER_DEPLOYED" = "1" ] || die "state does not record a deployed TiUP cluster"
PAUSE_WORKERS_STOPPED="${PAUSE_WORKERS_STOPPED:-0}"
PAUSE_TIUP_STOPPED="${PAUSE_TIUP_STOPPED:-0}"
CLUSTER_PAUSED="${CLUSTER_PAUSED:-0}"
for flag in PAUSE_WORKERS_STOPPED PAUSE_TIUP_STOPPED CLUSTER_PAUSED; do
  case "${!flag}" in 0|1) ;; *) die "invalid $flag value in state" ;; esac
done

need_cmd aws
account_id="$(aws --profile "$PROFILE" sts get-caller-identity --query Account --output text)"
[ "$account_id" = "$EXPECTED_AWS_ACCOUNT_ID" ] || \
  die "AWS account $account_id does not match expected $EXPECTED_AWS_ACCOUNT_ID"
for instance_id in $ALL_INSTANCE_IDS; do
  live_token="$(aws --profile "$PROFILE" --region "$REGION" ec2 describe-instances \
    --instance-ids "$instance_id" \
    --query "Reservations[0].Instances[0].Tags[?Key=='ProvisionToken'].Value | [0]" --output text)"
  live_cluster="$(aws --profile "$PROFILE" --region "$REGION" ec2 describe-instances \
    --instance-ids "$instance_id" \
    --query "Reservations[0].Instances[0].Tags[?Key=='ClusterName'].Value | [0]" --output text)"
  [ "$live_token" = "$PROVISION_TOKEN" ] && [ "$live_cluster" = "$CLUSTER_NAME" ] || \
    die "instance $instance_id does not belong to this recorded cluster"
done

if [ "$apply" -eq 0 ]; then
  cat <<EOF
PLAN ONLY: no component or instance will be stopped.
To pause, rerun: STATE_FILE="$STATE_FILE" $0 --apply

Will stop in order:
  cse-worker on every TiKV in cluster: $CLUSTER_NAME
  local TiUP cluster: $CLUSTER_NAME
  EC2 instances: $ALL_INSTANCE_IDS
  current phases: workers=$PAUSE_WORKERS_STOPPED tiup=$PAUSE_TIUP_STOPPED ec2=$CLUSTER_PAUSED
Will preserve:
  S3: s3://$S3_BUCKET/$S3_PREFIX/
  IAM role/profile: $TIKV_IAM_ROLE
  all EBS, security groups, and TiUP metadata
EOF
  exit 0
fi

need_cmd tiup

if [ "$PAUSE_WORKERS_STOPPED" = "0" ]; then
  run_step "stop workers" 600 \
    tiup cluster exec "$CLUSTER_NAME" -R tikv --command 'sudo systemctl stop cse-worker'
  write_state "$STATE_FILE" PAUSE_WORKERS_STOPPED "1"
fi
if [ "$PAUSE_TIUP_STOPPED" = "0" ]; then
  run_step "stop TiUP cluster" 900 tiup cluster stop "$CLUSTER_NAME" -y
  write_state "$STATE_FILE" PAUSE_TIUP_STOPPED "1"
fi
tiup cluster display "$CLUSTER_NAME" || true

if [ "$CLUSTER_PAUSED" = "0" ]; then
  aws --profile "$PROFILE" --region "$REGION" ec2 stop-instances \
    --instance-ids $ALL_INSTANCE_IDS
  run_step "wait for EC2 stopped" 1200 \
    aws --profile "$PROFILE" --region "$REGION" ec2 wait instance-stopped \
      --instance-ids $ALL_INSTANCE_IDS
  write_state "$STATE_FILE" CLUSTER_PAUSED "1"
fi

log "Cluster paused; S3, IAM, EBS, security groups, and TiUP metadata preserved"