#!/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:
  aws_provision.sh [--plan|--apply]

Creates the canonical TiDB X S3/IAM boundary, cluster security group, and EC2
nodes. TiKV nodes alone receive the S3 instance profile. Defaults to --plan.
USAGE
}

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

need_cmd aws
require_env \
  PROFILE EXPECTED_AWS_ACCOUNT_ID REGION AZ VPC_ID SUBNET_ID AMI KEY_NAME RUN_DIR RUN_LOG_DIR CLUSTER_NAME PROVISION_TOKEN OWNER \
  CREATED_BY CLEANUP_OWNER EXPIRES_AT COST_CEILING_USD CONTROL_HOST_ID CONTROL_SG OPERATOR_CIDR ASSOCIATE_PUBLIC_IP \
  ESTIMATED_BURN_RATE_USD_PER_HOUR ESTIMATED_BURN_RATE_USD_PER_DAY ESTIMATED_TOTAL_COST_USD CONTROL_TYPE NETWORK_EXTRAS \
  SSH_PORT TIDB_PORT GRAFANA_PORT PROMETHEUS_PORT \
  PD_COUNT PD_TYPE PD_ROOT_GB TIDB_COUNT TIDB_TYPE TIDB_ROOT_GB \
  TIKV_COUNT TIKV_TYPE TIKV_ROOT_GB MONITOR_COUNT MONITOR_TYPE \
  MONITOR_ROOT_GB S3_BUCKET S3_PREFIX TIKV_IAM_ROLE TIKV_IAM_POLICY \
  S3_NETWORK_MODE S3_NETWORK_RESOURCE_ID S3_NETWORK_POLICY_APPROVED

PRESSURE_COUNT="${PRESSURE_COUNT:-0}"
PRESSURE_TYPE="${PRESSURE_TYPE:-}"
PRESSURE_ROOT_GB="${PRESSURE_ROOT_GB:-}"
[ "$S3_NETWORK_POLICY_APPROVED" = "1" ] || \
  die "S3_NETWORK_POLICY_APPROVED must be 1 after reviewing the selected route and endpoint/network policy"
for name in PD_COUNT PD_ROOT_GB TIDB_COUNT TIDB_ROOT_GB TIKV_COUNT TIKV_ROOT_GB MONITOR_COUNT MONITOR_ROOT_GB; do
  require_positive_int "$name"
done
require_nonnegative_int PRESSURE_COUNT
if [ "$PRESSURE_COUNT" -gt 0 ]; then
  require_env PRESSURE_TYPE PRESSURE_ROOT_GB
  require_positive_int PRESSURE_ROOT_GB
fi
for name in COST_CEILING_USD ESTIMATED_BURN_RATE_USD_PER_HOUR ESTIMATED_BURN_RATE_USD_PER_DAY ESTIMATED_TOTAL_COST_USD; do
  require_nonnegative_decimal "$name"
done
awk -v estimate="$ESTIMATED_TOTAL_COST_USD" -v ceiling="$COST_CEILING_USD" \
  'BEGIN { exit !(estimate <= ceiling) }' || \
  die "estimated total cost exceeds the approved ceiling"
for name in SSH_PORT TIDB_PORT GRAFANA_PORT PROMETHEUS_PORT; do
  require_port "$name"
done

case "$S3_PREFIX" in
  ""|/*|*/|*//*|*[!A-Za-z0-9._/-]*) die "S3_PREFIX must be a safe non-empty relative path without trailing slash" ;;
esac
for name in CLUSTER_NAME TIKV_IAM_ROLE TIKV_IAM_POLICY OWNER CREATED_BY CLEANUP_OWNER; do require_safe_name "$name"; done
case "$S3_BUCKET" in ""|*[!a-z0-9.-]*) die "S3_BUCKET contains unsupported characters" ;; esac
case "$EXPIRES_AT" in ""|*[!0-9TtZz:+._-]*) die "EXPIRES_AT contains unsupported characters" ;; esac
case "$PROVISION_TOKEN" in
  ""|*[!A-Za-z0-9._-]*) die "PROVISION_TOKEN must contain only letters, digits, dot, underscore, or hyphen" ;;
esac
[ "${#PROVISION_TOKEN}" -le 48 ] || die "PROVISION_TOKEN must be at most 48 characters"

STATE_FILE="${STATE_FILE:-}"
require_env STATE_FILE
case "$RUN_DIR" in /*) ;; *) die "RUN_DIR must be an absolute path" ;; esac
case "$STATE_FILE" in "$RUN_DIR"/*) ;; *) die "STATE_FILE must be inside RUN_DIR" ;; esac
case "$RUN_LOG_DIR" in "$RUN_DIR"/*) ;; *) die "RUN_LOG_DIR must be inside RUN_DIR" ;; esac
validate_s3_network_selection "$S3_NETWORK_MODE" "$ASSOCIATE_PUBLIC_IP" "$S3_NETWORK_RESOURCE_ID"

if [ "$apply" -eq 0 ]; then
  cat <<EOF
PLAN ONLY: no AWS resources will be created.
To apply, rerun: STATE_FILE="$STATE_FILE" $0 --apply

Will create:
  provision token: $PROVISION_TOKEN
  private encrypted S3 bucket: $S3_BUCKET
  fresh S3 prefix: $S3_PREFIX
  TiKV role/profile: $TIKV_IAM_ROLE
  prefix policy: $TIKV_IAM_POLICY
  security group: $CLUSTER_NAME-sg in $VPC_ID
  nodes:
    pd=${PD_COUNT}x${PD_TYPE}/${PD_ROOT_GB}GiB
    tidb=${TIDB_COUNT}x${TIDB_TYPE}/${TIDB_ROOT_GB}GiB
    tikv=${TIKV_COUNT}x${TIKV_TYPE}/${TIKV_ROOT_GB}GiB with instance profile
    monitor=${MONITOR_COUNT}x${MONITOR_TYPE}/${MONITOR_ROOT_GB}GiB
    pressure=${PRESSURE_COUNT}x${PRESSURE_TYPE:-none}/${PRESSURE_ROOT_GB:-none}GiB
  public IP association: $ASSOCIATE_PUBLIC_IP
  operator ingress CIDR: $OPERATOR_CIDR
  operator ports: ssh=$SSH_PORT sql=$TIDB_PORT grafana=$GRAFANA_PORT prometheus=$PROMETHEUS_PORT
  S3 network path: $S3_NETWORK_MODE via $S3_NETWORK_RESOURCE_ID
Will preserve:
  control host: $CONTROL_HOST_ID via security group $CONTROL_SG
Launch context:
  AMI: $AMI
  key pair: $KEY_NAME
  cleanup owner: $CLEANUP_OWNER
State file:
  $STATE_FILE
EOF
  exit 0
fi

require_env APPROVED_PLAN_SHA256
actual_plan_sha256="$(aws_approval_sha256)"
[ "$APPROVED_PLAN_SHA256" = "$actual_plan_sha256" ] || \
  die "APPROVED_PLAN_SHA256 does not match this exact provision plan; rerun aws_plan.sh and obtain approval"

if [ -e "$STATE_FILE" ]; then
  die "state file already exists; review and clean its exact inventory or choose a new STATE_FILE"
fi

account_id="$(aws --profile "$PROFILE" sts get-caller-identity --query Account --output text)"
if [ "$account_id" != "$EXPECTED_AWS_ACCOUNT_ID" ]; then
  die "AWS account $account_id does not match expected $EXPECTED_AWS_ACCOUNT_ID"
fi
AMI_ROOT_DEVICE_NAME="$(resolve_ami_root_device "$PROFILE" "$REGION" "$AMI")"
network_route_table="$(validate_s3_network_path \
  "$PROFILE" "$REGION" "$VPC_ID" "$SUBNET_ID" \
  "$S3_NETWORK_MODE" "$ASSOCIATE_PUBLIC_IP" "$S3_NETWORK_RESOURCE_ID")"
log "Validated S3 network path through route table $network_route_table"
validate_subnet_az "$PROFILE" "$REGION" "$SUBNET_ID" "$AZ"

write_state "$STATE_FILE" PROFILE "$PROFILE"
write_state "$STATE_FILE" RUN_DIR "$RUN_DIR"
write_state "$STATE_FILE" RUN_LOG_DIR "$RUN_LOG_DIR"
write_state "$STATE_FILE" EXPECTED_AWS_ACCOUNT_ID "$EXPECTED_AWS_ACCOUNT_ID"
write_state "$STATE_FILE" REGION "$REGION"
write_state "$STATE_FILE" VPC_ID "$VPC_ID"
write_state "$STATE_FILE" SUBNET_ID "$SUBNET_ID"
write_state "$STATE_FILE" CLUSTER_NAME "$CLUSTER_NAME"
write_state "$STATE_FILE" PROVISION_TOKEN "$PROVISION_TOKEN"
write_state "$STATE_FILE" CONTROL_HOST_ID "$CONTROL_HOST_ID"
write_state "$STATE_FILE" CLEANUP_OWNER "$CLEANUP_OWNER"
write_state "$STATE_FILE" S3_NETWORK_MODE "$S3_NETWORK_MODE"
write_state "$STATE_FILE" S3_NETWORK_RESOURCE_ID "$S3_NETWORK_RESOURCE_ID"
write_state "$STATE_FILE" S3_NETWORK_POLICY_APPROVED "$S3_NETWORK_POLICY_APPROVED"
write_state "$STATE_FILE" S3_NETWORK_ROUTE_TABLE_ID "$network_route_table"
write_state "$STATE_FILE" SSH_PORT "$SSH_PORT"
write_state "$STATE_FILE" S3_BUCKET "$S3_BUCKET"
write_state "$STATE_FILE" S3_PREFIX "$S3_PREFIX"
write_state "$STATE_FILE" S3_BUCKET_CREATED "0"
write_state "$STATE_FILE" TIKV_IAM_ROLE "$TIKV_IAM_ROLE"
write_state "$STATE_FILE" TIKV_IAM_POLICY "$TIKV_IAM_POLICY"
write_state "$STATE_FILE" TIKV_IAM_PROFILE "$TIKV_IAM_ROLE"
write_state "$STATE_FILE" TIKV_IAM_PROFILE_ROLE_ATTACHED "0"
write_state "$STATE_FILE" CLUSTER_SG_NAME "$CLUSTER_NAME-sg"
write_state "$STATE_FILE" CLUSTER_SG ""
write_state "$STATE_FILE" TIUP_CLUSTER_DEPLOY_INTENT "0"
write_state "$STATE_FILE" TIUP_CLUSTER_DEPLOYED "0"
write_state "$STATE_FILE" PAUSE_WORKERS_STOPPED "0"
write_state "$STATE_FILE" PAUSE_TIUP_STOPPED "0"
write_state "$STATE_FILE" CLUSTER_PAUSED "0"
write_state "$STATE_FILE" ALL_INSTANCE_IDS ""
write_state "$STATE_FILE" ALL_VOLUME_IDS ""
write_state "$STATE_FILE" ALL_CLIENT_TOKENS ""

tag_spec() {
  local resource_type="$1"
  local name="$2"
  local extra_tags="${3:-}"
  printf 'ResourceType=%s,Tags=[{Key=Name,Value=%s},{Key=Owner,Value=%s},{Key=Purpose,Value=tidbx-aws-multinode-test},{Key=ClusterName,Value=%s},{Key=ProvisionToken,Value=%s},{Key=CreatedBy,Value=%s},{Key=ExpiresAt,Value=%s},{Key=CostCeilingUSD,Value=%s}%s]' \
    "$resource_type" "$name" "$OWNER" "$CLUSTER_NAME" "$PROVISION_TOKEN" "$CREATED_BY" "$EXPIRES_AT" "$COST_CEILING_USD" "$extra_tags"
}

log "Creating S3 bucket $S3_BUCKET"
if aws --profile "$PROFILE" --region "$REGION" s3api head-bucket \
  --bucket "$S3_BUCKET" >/dev/null 2>&1; then
  die "S3 bucket already exists in this account; choose a new globally unique bucket"
fi
if [ "$REGION" = "us-east-1" ]; then
  aws --profile "$PROFILE" --region "$REGION" s3api create-bucket \
    --bucket "$S3_BUCKET"
else
  aws --profile "$PROFILE" --region "$REGION" s3api create-bucket \
    --bucket "$S3_BUCKET" \
    --create-bucket-configuration "LocationConstraint=$REGION"
fi
aws --profile "$PROFILE" --region "$REGION" s3api put-public-access-block \
  --bucket "$S3_BUCKET" \
  --public-access-block-configuration \
  BlockPublicAcls=true,IgnorePublicAcls=true,BlockPublicPolicy=true,RestrictPublicBuckets=true
aws --profile "$PROFILE" --region "$REGION" s3api put-bucket-encryption \
  --bucket "$S3_BUCKET" \
  --server-side-encryption-configuration \
  '{"Rules":[{"ApplyServerSideEncryptionByDefault":{"SSEAlgorithm":"AES256"}}]}'
aws --profile "$PROFILE" --region "$REGION" s3api put-bucket-tagging \
  --bucket "$S3_BUCKET" \
  --tagging "TagSet=[{Key=Owner,Value=$OWNER},{Key=Purpose,Value=tidbx-aws-multinode-test},{Key=ClusterName,Value=$CLUSTER_NAME},{Key=ProvisionToken,Value=$PROVISION_TOKEN},{Key=CreatedBy,Value=$CREATED_BY},{Key=ExpiresAt,Value=$EXPIRES_AT},{Key=CostCeilingUSD,Value=$COST_CEILING_USD}]"
write_state "$STATE_FILE" S3_BUCKET_CREATED "1"

trust_policy='{"Version":"2012-10-17","Statement":[{"Effect":"Allow","Principal":{"Service":"ec2.amazonaws.com"},"Action":"sts:AssumeRole"}]}'
aws --profile "$PROFILE" iam create-role \
  --role-name "$TIKV_IAM_ROLE" \
  --assume-role-policy-document "$trust_policy" \
  --tags \
  "Key=Owner,Value=$OWNER" \
  "Key=Purpose,Value=tidbx-aws-multinode-test" \
  "Key=ClusterName,Value=$CLUSTER_NAME" \
  "Key=ProvisionToken,Value=$PROVISION_TOKEN" \
  "Key=CreatedBy,Value=$CREATED_BY" \
  "Key=ExpiresAt,Value=$EXPIRES_AT" \
  "Key=CostCeilingUSD,Value=$COST_CEILING_USD"

policy_document="$(printf '%s' "{
  \"Version\":\"2012-10-17\",
  \"Statement\":[
    {\"Effect\":\"Allow\",\"Action\":[\"s3:GetBucketLocation\"],\"Resource\":\"arn:aws:s3:::$S3_BUCKET\"},
    {\"Effect\":\"Allow\",\"Action\":[\"s3:ListBucket\"],\"Resource\":\"arn:aws:s3:::$S3_BUCKET\",\"Condition\":{\"StringLike\":{\"s3:prefix\":[\"$S3_PREFIX\",\"$S3_PREFIX/*\"]}}},
    {\"Effect\":\"Allow\",\"Action\":[\"s3:PutObject\",\"s3:GetObject\",\"s3:DeleteObject\",\"s3:AbortMultipartUpload\",\"s3:ListMultipartUploadParts\"],\"Resource\":\"arn:aws:s3:::$S3_BUCKET/$S3_PREFIX/*\"}
  ]
}")"
aws --profile "$PROFILE" iam put-role-policy \
  --role-name "$TIKV_IAM_ROLE" \
  --policy-name "$TIKV_IAM_POLICY" \
  --policy-document "$policy_document"
aws --profile "$PROFILE" iam create-instance-profile \
  --instance-profile-name "$TIKV_IAM_ROLE" \
  --tags \
  "Key=Owner,Value=$OWNER" \
  "Key=Purpose,Value=tidbx-aws-multinode-test" \
  "Key=ClusterName,Value=$CLUSTER_NAME" \
  "Key=ProvisionToken,Value=$PROVISION_TOKEN" \
  "Key=CreatedBy,Value=$CREATED_BY" \
  "Key=ExpiresAt,Value=$EXPIRES_AT" \
  "Key=CostCeilingUSD,Value=$COST_CEILING_USD"
aws --profile "$PROFILE" iam add-role-to-instance-profile \
  --instance-profile-name "$TIKV_IAM_ROLE" \
  --role-name "$TIKV_IAM_ROLE"
write_state "$STATE_FILE" TIKV_IAM_PROFILE_ROLE_ATTACHED "1"

profile_ready=0
for _ in $(seq 1 30); do
  if aws --profile "$PROFILE" iam get-instance-profile \
    --instance-profile-name "$TIKV_IAM_ROLE" \
    --query "InstanceProfile.Roles[?RoleName=='$TIKV_IAM_ROLE'].RoleName" \
    --output text | grep -qx "$TIKV_IAM_ROLE"; then
    profile_ready=1
    break
  fi
  sleep 2
done
[ "$profile_ready" -eq 1 ] || die "instance profile did not become ready"

preflight_file="$(mktemp)"
trap 'rm -f "$preflight_file"' EXIT
: > "$preflight_file"
aws --profile "$PROFILE" --region "$REGION" s3api put-object \
  --bucket "$S3_BUCKET" --key "$S3_PREFIX/.preflight" --body "$preflight_file"
aws --profile "$PROFILE" --region "$REGION" s3api head-object \
  --bucket "$S3_BUCKET" --key "$S3_PREFIX/.preflight"
aws --profile "$PROFILE" --region "$REGION" s3api delete-object \
  --bucket "$S3_BUCKET" --key "$S3_PREFIX/.preflight"

SG="$(aws --profile "$PROFILE" --region "$REGION" ec2 create-security-group \
  --group-name "$CLUSTER_NAME-sg" \
  --description "$CLUSTER_NAME TiDB X test cluster" \
  --vpc-id "$VPC_ID" \
  --tag-specifications "$(tag_spec security-group "$CLUSTER_NAME-sg")" \
  --query GroupId --output text)"
write_state "$STATE_FILE" SG "$SG"
write_state "$STATE_FILE" CLUSTER_SG "$SG"
write_state "$STATE_FILE" CONTROL_SG "$CONTROL_SG"

aws --profile "$PROFILE" --region "$REGION" ec2 authorize-security-group-ingress \
  --group-id "$SG" --source-group "$SG" --protocol -1 --port -1
aws --profile "$PROFILE" --region "$REGION" ec2 authorize-security-group-ingress \
  --group-id "$SG" --source-group "$CONTROL_SG" --protocol -1 --port -1
aws --profile "$PROFILE" --region "$REGION" ec2 authorize-security-group-ingress \
  --group-id "$SG" --ip-permissions \
  "IpProtocol=tcp,FromPort=$SSH_PORT,ToPort=$SSH_PORT,IpRanges=[{CidrIp=$OPERATOR_CIDR,Description=operator-ssh}]" \
  "IpProtocol=tcp,FromPort=$TIDB_PORT,ToPort=$TIDB_PORT,IpRanges=[{CidrIp=$OPERATOR_CIDR,Description=operator-sql}]" \
  "IpProtocol=tcp,FromPort=$GRAFANA_PORT,ToPort=$GRAFANA_PORT,IpRanges=[{CidrIp=$OPERATOR_CIDR,Description=operator-grafana}]" \
  "IpProtocol=tcp,FromPort=$PROMETHEUS_PORT,ToPort=$PROMETHEUS_PORT,IpRanges=[{CidrIp=$OPERATOR_CIDR,Description=operator-prometheus}]"

network_args=()
if [ "$ASSOCIATE_PUBLIC_IP" = "1" ]; then
  network_args+=(--associate-public-ip-address)
else
  network_args+=(--no-associate-public-ip-address)
fi

launch_node() {
  local role="$1"
  local idx="$2"
  local instance_type="$3"
  local root_gb="$4"
  local client_token="$5"
  local name="$CLUSTER_NAME-$role-$idx"
  local extra=",{Key=Role,Value=$role},{Key=Index,Value=$idx}"

  aws --profile "$PROFILE" --region "$REGION" ec2 run-instances \
    --client-token "$client_token" \
    --image-id "$AMI" \
    --instance-type "$instance_type" \
    --key-name "$KEY_NAME" \
    --subnet-id "$SUBNET_ID" \
    --security-group-ids "$SG" \
    "${network_args[@]}" \
    --block-device-mappings "[{\"DeviceName\":\"$AMI_ROOT_DEVICE_NAME\",\"Ebs\":{\"VolumeSize\":$root_gb,\"VolumeType\":\"gp3\",\"DeleteOnTermination\":true}}]" \
    --tag-specifications "$(tag_spec instance "$name" "$extra")" \
    --query 'Instances[0].InstanceId' --output text
}

launch_tikv() {
  local idx="$1"
  local client_token="$2"
  local name="$CLUSTER_NAME-tikv-$idx"
  local extra=",{Key=Role,Value=tikv},{Key=Index,Value=$idx}"

  aws --profile "$PROFILE" --region "$REGION" ec2 run-instances \
    --client-token "$client_token" \
    --image-id "$AMI" \
    --instance-type "$TIKV_TYPE" \
    --key-name "$KEY_NAME" \
    --subnet-id "$SUBNET_ID" \
    --security-group-ids "$SG" \
    "${network_args[@]}" \
    --iam-instance-profile "Name=$TIKV_IAM_ROLE" \
    --block-device-mappings "[{\"DeviceName\":\"$AMI_ROOT_DEVICE_NAME\",\"Ebs\":{\"VolumeSize\":$TIKV_ROOT_GB,\"VolumeType\":\"gp3\",\"DeleteOnTermination\":true}}]" \
    --tag-specifications "$(tag_spec instance "$name" "$extra")" \
    --query 'Instances[0].InstanceId' --output text
}

all_instance_ids=()
all_volume_ids=()
all_client_tokens=()
group_ids=()

record_instance() {
  local role="$1"
  local id="$2"
  local group_state_key
  local volume_ids
  local volume_id
  local attempt

  case "$role" in
    pd) group_state_key="PD_IDS" ;;
    tidb) group_state_key="TIDB_IDS" ;;
    tikv) group_state_key="TIKV_IDS" ;;
    monitor) group_state_key="MONITOR_IDS" ;;
    pressure) group_state_key="PRESSURE_IDS" ;;
    *) die "unsupported instance role: $role" ;;
  esac

  all_instance_ids+=("$id")
  group_ids+=("$id")
  write_state "$STATE_FILE" "$group_state_key" "${group_ids[*]}"
  write_state "$STATE_FILE" ALL_INSTANCE_IDS "${all_instance_ids[*]}"

  volume_ids=""
  for attempt in $(seq 1 15); do
    volume_ids="$(aws --profile "$PROFILE" --region "$REGION" ec2 describe-instances \
      --instance-ids "$id" \
      --query 'Reservations[0].Instances[0].BlockDeviceMappings[].Ebs.VolumeId' \
      --output text 2>/dev/null || true)"
    [ -n "$volume_ids" ] && [ "$volume_ids" != "None" ] && break
    sleep 1
  done
  [ -n "$volume_ids" ] && [ "$volume_ids" != "None" ] || \
    die "instance $id was recorded but its EBS inventory could not be resolved"
  for volume_id in $volume_ids; do
    all_volume_ids+=("$volume_id")
  done
  write_state "$STATE_FILE" ALL_VOLUME_IDS "${all_volume_ids[*]}"
}

append_group() {
  local role="$1"
  local count="$2"
  local instance_type="$3"
  local root_gb="$4"
  local id
  local i
  local client_token

  group_ids=()
  for i in $(seq 1 "$count"); do
    client_token="$PROVISION_TOKEN-$role-$i"
    all_client_tokens+=("$client_token")
    write_state "$STATE_FILE" ALL_CLIENT_TOKENS "${all_client_tokens[*]}"
    id="$(launch_node "$role" "$i" "$instance_type" "$root_gb" "$client_token")"
    record_instance "$role" "$id"
  done
}

append_group pd "$PD_COUNT" "$PD_TYPE" "$PD_ROOT_GB"
pd_ids="${group_ids[*]}"
write_state "$STATE_FILE" PD_IDS "$pd_ids"
append_group tidb "$TIDB_COUNT" "$TIDB_TYPE" "$TIDB_ROOT_GB"
tidb_ids="${group_ids[*]}"
write_state "$STATE_FILE" TIDB_IDS "$tidb_ids"

tikv_ids=()
group_ids=()
for i in $(seq 1 "$TIKV_COUNT"); do
  client_token="$PROVISION_TOKEN-tikv-$i"
  all_client_tokens+=("$client_token")
  write_state "$STATE_FILE" ALL_CLIENT_TOKENS "${all_client_tokens[*]}"
  id="$(launch_tikv "$i" "$client_token")"
  record_instance tikv "$id"
  tikv_ids+=("$id")
done

append_group monitor "$MONITOR_COUNT" "$MONITOR_TYPE" "$MONITOR_ROOT_GB"
monitor_ids="${group_ids[*]}"
write_state "$STATE_FILE" MONITOR_IDS "$monitor_ids"

if [ "$PRESSURE_COUNT" -gt 0 ]; then
  append_group pressure "$PRESSURE_COUNT" "$PRESSURE_TYPE" "$PRESSURE_ROOT_GB"
  pressure_ids="${group_ids[*]}"
  write_state "$STATE_FILE" PRESSURE_IDS "$pressure_ids"
fi

run_step "wait for EC2 running" 1200 \
  aws --profile "$PROFILE" --region "$REGION" ec2 wait instance-running \
    --instance-ids "${all_instance_ids[@]}"
run_step "wait for EC2 status checks" 1200 \
  aws --profile "$PROFILE" --region "$REGION" ec2 wait instance-status-ok \
    --instance-ids "${all_instance_ids[@]}"

verified_volume_ids="$(aws --profile "$PROFILE" --region "$REGION" ec2 describe-instances \
  --instance-ids "${all_instance_ids[@]}" \
  --query 'Reservations[].Instances[].BlockDeviceMappings[].Ebs.VolumeId' \
  --output text)"
write_state "$STATE_FILE" ALL_VOLUME_IDS "$verified_volume_ids"

private_ips_for_ids() {
  aws --profile "$PROFILE" --region "$REGION" ec2 describe-instances \
    --instance-ids $1 \
    --query 'Reservations[].Instances[].PrivateIpAddress' --output text
}

pd_hosts="$(private_ips_for_ids "$pd_ids")"
tidb_hosts="$(private_ips_for_ids "$tidb_ids")"
tikv_hosts="$(private_ips_for_ids "${tikv_ids[*]}")"
monitor_hosts="$(private_ips_for_ids "$monitor_ids")"
write_state "$STATE_FILE" PD_HOSTS "$pd_hosts"
write_state "$STATE_FILE" TIDB_HOSTS "$tidb_hosts"
write_state "$STATE_FILE" TIKV_HOSTS "$tikv_hosts"
write_state "$STATE_FILE" MONITOR_HOSTS "$monitor_hosts"
if [ "$PRESSURE_COUNT" -gt 0 ]; then
  pressure_hosts="$(private_ips_for_ids "$pressure_ids")"
  write_state "$STATE_FILE" PRESSURE_HOSTS "$pressure_hosts"
fi

profile_arns="$(aws --profile "$PROFILE" --region "$REGION" ec2 describe-instances \
  --instance-ids "${tikv_ids[@]}" \
  --query 'Reservations[].Instances[].IamInstanceProfile.Arn' --output text)"
profile_count=0
for profile_arn in $profile_arns; do
  case "$profile_arn" in
    */"$TIKV_IAM_ROLE") profile_count=$((profile_count + 1)) ;;
    *) die "unexpected TiKV instance profile: $profile_arn" ;;
  esac
done
[ "$profile_count" -eq "$TIKV_COUNT" ] || die "not every TiKV has the expected instance profile"

aws --profile "$PROFILE" --region "$REGION" ec2 describe-instances \
  --instance-ids "${all_instance_ids[@]}" \
  --query 'Reservations[].Instances[].{Name:Tags[?Key==`Name`]|[0].Value,InstanceId:InstanceId,PrivateIp:PrivateIpAddress,PublicIp:PublicIpAddress,Profile:IamInstanceProfile.Arn,State:State.Name}' \
  --output table
printf 'Root EBS volumes: %s\n' "$verified_volume_ids"

log "Provisioning complete; state file: $STATE_FILE"