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

Runs on the control/build host. Destroys the local TiUP cluster when present and
the exact AWS/S3/IAM inventory recorded before a successful or interrupted
provision. Defaults to plan and requires a separate --apply.
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 \
  ALL_VOLUME_IDS ALL_CLIENT_TOKENS S3_BUCKET S3_PREFIX S3_BUCKET_CREATED \
  TIKV_IAM_ROLE TIKV_IAM_PROFILE TIKV_IAM_PROFILE_ROLE_ATTACHED \
  TIKV_IAM_POLICY CLUSTER_SG CLUSTER_SG_NAME TIUP_CLUSTER_DEPLOYED TIUP_CLUSTER_DEPLOY_INTENT CLEANUP_OWNER RUN_LOG_DIR \
  PROVISION_TOKEN VPC_ID
load_state_file "$STATE_FILE"
require_env \
  PROFILE EXPECTED_AWS_ACCOUNT_ID REGION CLUSTER_NAME CLEANUP_OWNER RUN_LOG_DIR PROVISION_TOKEN VPC_ID

ALL_INSTANCE_IDS="${ALL_INSTANCE_IDS:-}"
ALL_VOLUME_IDS="${ALL_VOLUME_IDS:-}"
ALL_CLIENT_TOKENS="${ALL_CLIENT_TOKENS:-}"
S3_BUCKET="${S3_BUCKET:-}"
S3_PREFIX="${S3_PREFIX:-}"
S3_BUCKET_CREATED="${S3_BUCKET_CREATED:-0}"
TIKV_IAM_ROLE="${TIKV_IAM_ROLE:-}"
TIKV_IAM_PROFILE="${TIKV_IAM_PROFILE:-}"
TIKV_IAM_PROFILE_ROLE_ATTACHED="${TIKV_IAM_PROFILE_ROLE_ATTACHED:-0}"
TIKV_IAM_POLICY="${TIKV_IAM_POLICY:-}"
CLUSTER_SG="${CLUSTER_SG:-}"
CLUSTER_SG_NAME="${CLUSTER_SG_NAME:-$CLUSTER_NAME-sg}"
TIUP_CLUSTER_DEPLOYED="${TIUP_CLUSTER_DEPLOYED:-0}"
TIUP_CLUSTER_DEPLOY_INTENT="${TIUP_CLUSTER_DEPLOY_INTENT:-0}"
PRESERVE_S3_BUCKET="${PRESERVE_S3_BUCKET:-0}"
case "$PRESERVE_S3_BUCKET" in 0|1) ;; *) die "PRESERVE_S3_BUCKET must be 0 or 1" ;; esac
case "$S3_BUCKET_CREATED" in 0|1) ;; *) die "invalid S3_BUCKET_CREATED value in state" ;; esac
case "$TIUP_CLUSTER_DEPLOYED" in 0|1) ;; *) die "invalid TIUP_CLUSTER_DEPLOYED value in state" ;; esac
case "$TIUP_CLUSTER_DEPLOY_INTENT" in 0|1) ;; *) die "invalid TIUP_CLUSTER_DEPLOY_INTENT value in state" ;; esac
case "$TIKV_IAM_PROFILE_ROLE_ATTACHED" in 0|1) ;; *) die "invalid IAM attachment value in state" ;; esac
if [ -n "$S3_BUCKET" ] || [ -n "$S3_PREFIX" ]; then
  [ -n "$S3_BUCKET" ] && [ -n "$S3_PREFIX" ] || \
    die "state must record both S3_BUCKET and S3_PREFIX"
fi
[ "$S3_BUCKET_CREATED" = "0" ] || [ -n "$S3_BUCKET" ] || \
  die "state proves bucket creation but records no bucket"

display_or_none() {
  if [ -n "$1" ]; then
    printf '%s' "$1"
  else
    printf '<none recorded>'
  fi
}

append_unique_id() {
  local id="$1"
  case " $ALL_INSTANCE_IDS " in
    *" $id "*) ;;
    *) ALL_INSTANCE_IDS="${ALL_INSTANCE_IDS:+$ALL_INSTANCE_IDS }$id" ;;
  esac
}

require_live_tag() {
  local service="$1"
  local resource_id="$2"
  local key="$3"
  local expected="$4"
  local actual

  case "$service" in
    instance)
      actual="$(aws --profile "$PROFILE" --region "$REGION" ec2 describe-instances \
        --instance-ids "$resource_id" \
        --query "Reservations[0].Instances[0].Tags[?Key=='$key'].Value | [0]" --output text)"
      ;;
    security-group)
      actual="$(aws --profile "$PROFILE" --region "$REGION" ec2 describe-security-groups \
        --group-ids "$resource_id" \
        --query "SecurityGroups[0].Tags[?Key=='$key'].Value | [0]" --output text)"
      ;;
    role)
      actual="$(aws --profile "$PROFILE" iam list-role-tags --role-name "$resource_id" \
        --query "Tags[?Key=='$key'].Value | [0]" --output text)"
      ;;
    instance-profile)
      actual="$(aws --profile "$PROFILE" iam list-instance-profile-tags \
        --instance-profile-name "$resource_id" \
        --query "Tags[?Key=='$key'].Value | [0]" --output text)"
      ;;
    bucket)
      actual="$(aws --profile "$PROFILE" --region "$REGION" s3api get-bucket-tagging \
        --bucket "$resource_id" --query "TagSet[?Key=='$key'].Value | [0]" --output text)"
      ;;
    *) die "unsupported tag service: $service" ;;
  esac
  [ "$actual" = "$expected" ] || \
    die "$service $resource_id has $key=${actual:-<missing>}, expected $expected"
}

resource_name_exists() {
  local kind="$1"
  local name="$2"
  local result
  case "$kind" in
    bucket)
      result="$(aws --profile "$PROFILE" s3api list-buckets \
        --query "Buckets[?Name=='$name'].Name | [0]" --output text)" || \
        die "failed to inventory S3 bucket $name"
      ;;
    role)
      result="$(aws --profile "$PROFILE" iam list-roles \
        --query "Roles[?RoleName=='$name'].RoleName | [0]" --output text)" || \
        die "failed to inventory IAM role $name"
      ;;
    instance-profile)
      result="$(aws --profile "$PROFILE" iam list-instance-profiles \
        --query "InstanceProfiles[?InstanceProfileName=='$name'].InstanceProfileName | [0]" --output text)" || \
        die "failed to inventory IAM instance profile $name"
      ;;
    *) die "unsupported resource kind: $kind" ;;
  esac
  [ -n "$result" ] && [ "$result" != "None" ]
}

cleanup_manifest() {
  printf '%s\n' \
    "STATE_FILE=$STATE_FILE" \
    "EXPECTED_AWS_ACCOUNT_ID=$EXPECTED_AWS_ACCOUNT_ID" \
    "REGION=$REGION" \
    "CLUSTER_NAME=$CLUSTER_NAME" \
    "PROVISION_TOKEN=$PROVISION_TOKEN" \
    "TIUP_CLUSTER_DEPLOYED=$TIUP_CLUSTER_DEPLOYED" \
    "TIUP_CLUSTER_DEPLOY_INTENT=$TIUP_CLUSTER_DEPLOY_INTENT" \
    "ALL_INSTANCE_IDS=$ALL_INSTANCE_IDS" \
    "ALL_VOLUME_IDS=$ALL_VOLUME_IDS" \
    "S3_BUCKET=$S3_BUCKET" \
    "S3_PREFIX=$S3_PREFIX" \
    "S3_BUCKET_CREATED=$S3_BUCKET_CREATED" \
    "PRESERVE_S3_BUCKET=$PRESERVE_S3_BUCKET" \
    "TIKV_IAM_ROLE=$TIKV_IAM_ROLE" \
    "TIKV_IAM_PROFILE=$TIKV_IAM_PROFILE" \
    "TIKV_IAM_POLICY=$TIKV_IAM_POLICY" \
    "CLUSTER_SG=$CLUSTER_SG"
}

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"

if [ -z "$CLUSTER_SG" ]; then
  recovered_sg="$(aws --profile "$PROFILE" --region "$REGION" ec2 describe-security-groups \
    --filters "Name=vpc-id,Values=$VPC_ID" "Name=group-name,Values=$CLUSTER_SG_NAME" \
      "Name=tag:ProvisionToken,Values=$PROVISION_TOKEN" \
    --query 'SecurityGroups[].GroupId' --output text)"
  if [ -n "$recovered_sg" ] && [ "$recovered_sg" != "None" ]; then
    set -- $recovered_sg
    [ "$#" -eq 1 ] || die "multiple security groups match the recorded creation intent"
    CLUSTER_SG="$1"
    write_state "$STATE_FILE" CLUSTER_SG "$CLUSTER_SG"
  fi
fi

for client_token in $ALL_CLIENT_TOKENS; do
  recovered_ids=""
  for attempt in $(seq 1 6); do
    recovered_ids="$(aws --profile "$PROFILE" --region "$REGION" ec2 describe-instances \
      --filters "Name=client-token,Values=$client_token" \
      --query 'Reservations[].Instances[].InstanceId' --output text)"
    if [ -n "$recovered_ids" ] && [ "$recovered_ids" != "None" ]; then
      break
    fi
    [ "$attempt" -eq 6 ] || sleep 5
  done
  if [ -z "$recovered_ids" ] || [ "$recovered_ids" = "None" ]; then
    log "No EC2 instance is visible for recorded client token $client_token"
    continue
  fi
  for recovered_id in $recovered_ids; do
    append_unique_id "$recovered_id"
  done
done
write_state "$STATE_FILE" ALL_INSTANCE_IDS "$ALL_INSTANCE_IDS"

if [ -n "$ALL_INSTANCE_IDS" ]; then
  instance_id_csv="${ALL_INSTANCE_IDS// /,}"
  live_instance_ids="$(aws --profile "$PROFILE" --region "$REGION" ec2 describe-instances \
    --filters "Name=instance-id,Values=$instance_id_csv" \
    --query 'Reservations[].Instances[].InstanceId' --output text)"
  if [ "$live_instance_ids" = "None" ]; then
    live_instance_ids=""
  fi
  ALL_INSTANCE_IDS="$live_instance_ids"
  write_state "$STATE_FILE" ALL_INSTANCE_IDS "$ALL_INSTANCE_IDS"

  discovered_volume_ids="$(aws --profile "$PROFILE" --region "$REGION" ec2 describe-instances \
    --filters "Name=instance-id,Values=$instance_id_csv" \
    --query 'Reservations[].Instances[].BlockDeviceMappings[].Ebs.VolumeId' --output text)"
  if [ -n "$discovered_volume_ids" ] && [ "$discovered_volume_ids" != "None" ]; then
    ALL_VOLUME_IDS="$discovered_volume_ids"
    write_state "$STATE_FILE" ALL_VOLUME_IDS "$ALL_VOLUME_IDS"
  elif [ -z "$ALL_INSTANCE_IDS" ]; then
    ALL_VOLUME_IDS=""
    write_state "$STATE_FILE" ALL_VOLUME_IDS ""
  fi
fi

for instance_id in $ALL_INSTANCE_IDS; do
  require_live_tag instance "$instance_id" ProvisionToken "$PROVISION_TOKEN"
  require_live_tag instance "$instance_id" ClusterName "$CLUSTER_NAME"
done
if [ -n "$S3_BUCKET" ] && [ "$S3_BUCKET_CREATED" = "1" ] && \
  resource_name_exists bucket "$S3_BUCKET"; then
  require_live_tag bucket "$S3_BUCKET" ProvisionToken "$PROVISION_TOKEN"
  require_live_tag bucket "$S3_BUCKET" ClusterName "$CLUSTER_NAME"
fi
if [ -n "$TIKV_IAM_ROLE" ] && resource_name_exists role "$TIKV_IAM_ROLE"; then
  require_live_tag role "$TIKV_IAM_ROLE" ProvisionToken "$PROVISION_TOKEN"
  require_live_tag role "$TIKV_IAM_ROLE" ClusterName "$CLUSTER_NAME"
fi
if [ -n "$TIKV_IAM_PROFILE" ] && resource_name_exists instance-profile "$TIKV_IAM_PROFILE"; then
  require_live_tag instance-profile "$TIKV_IAM_PROFILE" ProvisionToken "$PROVISION_TOKEN"
  require_live_tag instance-profile "$TIKV_IAM_PROFILE" ClusterName "$CLUSTER_NAME"
fi
if [ -n "$CLUSTER_SG" ]; then
  require_live_tag security-group "$CLUSTER_SG" ProvisionToken "$PROVISION_TOKEN"
  require_live_tag security-group "$CLUSTER_SG" ClusterName "$CLUSTER_NAME"
fi

cleanup_sha256="$(cleanup_manifest | sha256_stream)"

if [ "$apply" -eq 0 ]; then
  cat <<EOF
PLAN ONLY: no resource will be deleted.
To destroy, rerun after exact inventory review:
  STATE_FILE="$STATE_FILE" $0 --apply

Recorded cleanup inventory or creation intent:
  cleanup owner: $CLEANUP_OWNER
  local TiUP cluster: $CLUSTER_NAME (intent=$TIUP_CLUSTER_DEPLOY_INTENT, deployed=$TIUP_CLUSTER_DEPLOYED)
  EC2 instances and delete-on-termination root EBS: $(display_or_none "$ALL_INSTANCE_IDS")
  recorded root EBS volume IDs: $(display_or_none "$ALL_VOLUME_IDS")
  S3 bucket: $(display_or_none "$S3_BUCKET") (creation-proven=$S3_BUCKET_CREATED, preserve=$PRESERVE_S3_BUCKET)
  S3 prefix: $(display_or_none "$S3_PREFIX")
  IAM role: $(display_or_none "$TIKV_IAM_ROLE")
  IAM instance profile: $(display_or_none "$TIKV_IAM_PROFILE")
  IAM inline policy: $(display_or_none "$TIKV_IAM_POLICY")
  cluster security group: $(display_or_none "$CLUSTER_SG")
Approval SHA256: $cleanup_sha256
Will preserve:
  existing control/build host and every resource absent from the state file
  any S3 creation intent whose successful creation and provision-token tags cannot be proven
EOF
  exit 0
fi

require_env APPROVED_CLEANUP_SHA256
[ "$APPROVED_CLEANUP_SHA256" = "$cleanup_sha256" ] || \
  die "APPROVED_CLEANUP_SHA256 does not match this exact reconciled inventory; rerun cleanup plan"

if [ "$TIUP_CLUSTER_DEPLOYED" = "1" ] || [ "$TIUP_CLUSTER_DEPLOY_INTENT" = "1" ]; then
  need_cmd tiup
  tiup_clusters="$(tiup cluster list)" || die "could not list local TiUP clusters"
  if printf '%s\n' "$tiup_clusters" | awk 'NR > 1 {print $1}' | grep -Fxq "$CLUSTER_NAME"; then
    tiup cluster exec "$CLUSTER_NAME" -R tikv \
      --command 'sudo systemctl stop cse-worker' || true
    run_step "destroy TiUP cluster" 1800 \
      tiup cluster destroy "$CLUSTER_NAME" -y
  else
    log "TiUP cluster is already absent; reconciling state"
  fi
  write_state "$STATE_FILE" TIUP_CLUSTER_DEPLOYED "0"
  write_state "$STATE_FILE" TIUP_CLUSTER_DEPLOY_INTENT "0"
else
  log "No deployed TiUP cluster is recorded; skipping local TiUP cleanup"
fi

if [ -n "$ALL_INSTANCE_IDS" ]; then
  aws --profile "$PROFILE" --region "$REGION" ec2 terminate-instances \
    --instance-ids $ALL_INSTANCE_IDS
  run_step "wait for EC2 terminated" 1200 \
    aws --profile "$PROFILE" --region "$REGION" ec2 wait instance-terminated \
      --instance-ids $ALL_INSTANCE_IDS
fi
for volume_id in $ALL_VOLUME_IDS; do
  volume_count="$(aws --profile "$PROFILE" --region "$REGION" ec2 describe-volumes \
    --filters "Name=volume-id,Values=$volume_id" --query 'length(Volumes)' --output text)"
  [ "$volume_count" -eq 0 ] || die "root EBS volume still exists after termination: $volume_id"
done
if [ -n "$ALL_INSTANCE_IDS" ]; then
  write_state "$STATE_FILE" ALL_INSTANCE_IDS ""
  write_state "$STATE_FILE" ALL_VOLUME_IDS ""
  write_state "$STATE_FILE" ALL_CLIENT_TOKENS ""
fi

unresolved_cleanup=0
if [ -n "$S3_BUCKET" ]; then
  if ! resource_name_exists bucket "$S3_BUCKET"; then
    log "S3 bucket is already absent; reconciling state"
    write_state "$STATE_FILE" S3_BUCKET ""
    write_state "$STATE_FILE" S3_PREFIX ""
    write_state "$STATE_FILE" S3_BUCKET_CREATED "0"
  elif [ "$S3_BUCKET_CREATED" != "1" ]; then
    log "Preserving unverified bucket creation intent $S3_BUCKET; audit it manually"
    unresolved_cleanup=1
  else
    require_live_tag bucket "$S3_BUCKET" ProvisionToken "$PROVISION_TOKEN"
    require_live_tag bucket "$S3_BUCKET" ClusterName "$CLUSTER_NAME"
    run_step "delete S3 prefix" 1800 \
      aws --profile "$PROFILE" --region "$REGION" s3 rm \
        "s3://$S3_BUCKET/$S3_PREFIX/" --recursive
    remaining="$(aws --profile "$PROFILE" --region "$REGION" s3api list-objects-v2 \
      --bucket "$S3_BUCKET" --prefix "$S3_PREFIX/" --max-keys 1 \
      --query KeyCount --output text)"
    [ "$remaining" -eq 0 ] || die "S3 prefix is not empty after cleanup"

    if [ "$PRESERVE_S3_BUCKET" != "1" ]; then
      bucket_objects="$(aws --profile "$PROFILE" --region "$REGION" s3api list-objects-v2 \
        --bucket "$S3_BUCKET" --max-keys 1 --query KeyCount --output text)"
      [ "$bucket_objects" -eq 0 ] || die "bucket contains objects outside the recorded prefix"
      aws --profile "$PROFILE" --region "$REGION" s3api delete-bucket --bucket "$S3_BUCKET"
    else
      write_state "$STATE_FILE" PRESERVED_S3_BUCKET "$S3_BUCKET"
    fi
    write_state "$STATE_FILE" S3_BUCKET ""
    write_state "$STATE_FILE" S3_PREFIX ""
    write_state "$STATE_FILE" S3_BUCKET_CREATED "0"
  fi
fi

if [ -n "$TIKV_IAM_PROFILE" ]; then
  if resource_name_exists instance-profile "$TIKV_IAM_PROFILE"; then
    require_live_tag instance-profile "$TIKV_IAM_PROFILE" ProvisionToken "$PROVISION_TOKEN"
    require_live_tag instance-profile "$TIKV_IAM_PROFILE" ClusterName "$CLUSTER_NAME"
    attached_role="$(aws --profile "$PROFILE" iam get-instance-profile \
      --instance-profile-name "$TIKV_IAM_PROFILE" \
      --query "InstanceProfile.Roles[?RoleName=='$TIKV_IAM_ROLE'].RoleName | [0]" --output text)"
    if [ "$attached_role" = "$TIKV_IAM_ROLE" ]; then
      aws --profile "$PROFILE" iam remove-role-from-instance-profile \
        --instance-profile-name "$TIKV_IAM_PROFILE" --role-name "$TIKV_IAM_ROLE"
    fi
    aws --profile "$PROFILE" iam delete-instance-profile \
      --instance-profile-name "$TIKV_IAM_PROFILE"
  fi
  write_state "$STATE_FILE" TIKV_IAM_PROFILE_ROLE_ATTACHED "0"
  write_state "$STATE_FILE" TIKV_IAM_PROFILE ""
fi
if [ -n "$TIKV_IAM_POLICY" ]; then
  if resource_name_exists role "$TIKV_IAM_ROLE"; then
    require_live_tag role "$TIKV_IAM_ROLE" ProvisionToken "$PROVISION_TOKEN"
    policies="$(aws --profile "$PROFILE" iam list-role-policies --role-name "$TIKV_IAM_ROLE" \
      --query 'PolicyNames' --output text)"
    case " $policies " in
      *" $TIKV_IAM_POLICY "*) aws --profile "$PROFILE" iam delete-role-policy \
        --role-name "$TIKV_IAM_ROLE" --policy-name "$TIKV_IAM_POLICY" ;;
    esac
  fi
  write_state "$STATE_FILE" TIKV_IAM_POLICY ""
fi
if [ -n "$TIKV_IAM_ROLE" ]; then
  if resource_name_exists role "$TIKV_IAM_ROLE"; then
    require_live_tag role "$TIKV_IAM_ROLE" ProvisionToken "$PROVISION_TOKEN"
    require_live_tag role "$TIKV_IAM_ROLE" ClusterName "$CLUSTER_NAME"
    aws --profile "$PROFILE" iam delete-role --role-name "$TIKV_IAM_ROLE"
  fi
  write_state "$STATE_FILE" TIKV_IAM_ROLE ""
fi
if [ -n "$CLUSTER_SG" ]; then
  sg_count="$(aws --profile "$PROFILE" --region "$REGION" ec2 describe-security-groups \
    --filters "Name=group-id,Values=$CLUSTER_SG" --query 'length(SecurityGroups)' --output text)"
  if [ "$sg_count" -gt 0 ]; then
    require_live_tag security-group "$CLUSTER_SG" ProvisionToken "$PROVISION_TOKEN"
    require_live_tag security-group "$CLUSTER_SG" ClusterName "$CLUSTER_NAME"
    aws --profile "$PROFILE" --region "$REGION" ec2 delete-security-group --group-id "$CLUSTER_SG"
  fi
  write_state "$STATE_FILE" CLUSTER_SG ""
fi

if [ "$unresolved_cleanup" -ne 0 ]; then
  die "cleanup completed for verified resources but an unverified S3 creation intent remains"
fi
log "Cleanup complete for the exact recorded TiDB X inventory"