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

Validates the canonical TiDB X runtime, every worker, TiKV instance profiles,
real-S3 configuration and writes, SQL smoke, and ANALYZE. 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 PD_IDS PD_HOSTS \
  TIDB_IDS TIDB_HOSTS TIKV_IDS TIKV_HOSTS S3_BUCKET S3_PREFIX TIKV_IAM_ROLE TIKV_IAM_POLICY \
  TIUP_CLUSTER_DEPLOYED PROVISION_TOKEN TIDB_PORT TIDB_STATUS_PORT TIKV_PORT \
  TIKV_STATUS_PORT TOPOLOGY_FILE WORKER_CONFIG TIDB_RUNTIME_USER \
  TIDB_RUNTIME_GROUP TIDB_RUNTIME_HOME TIUP_DEPLOY_DIR KEYSPACE_NAME
load_state_file "$STATE_FILE"

MYSQL_USER="${MYSQL_USER:-}"
MYSQL_PASSWORD="${MYSQL_PASSWORD:-}"
WORKER_LOG="$TIDB_RUNTIME_HOME/cse-worker.log"
TIKV_LOG="$TIUP_DEPLOY_DIR/tikv-$TIKV_PORT/log/tikv.log"

require_env \
  PROFILE EXPECTED_AWS_ACCOUNT_ID REGION CLUSTER_NAME PD_IDS PD_HOSTS TIDB_IDS TIDB_HOSTS \
  TIUP_CLUSTER_DEPLOYED PROVISION_TOKEN TIDB_PORT TIDB_STATUS_PORT TIKV_HOSTS TIKV_IDS \
  TIKV_PORT TIKV_STATUS_PORT KEYSPACE_NAME \
  S3_BUCKET S3_PREFIX TIKV_IAM_ROLE TIKV_IAM_POLICY MYSQL_USER TOPOLOGY_FILE WORKER_CONFIG \
  TIDB_RUNTIME_USER TIDB_RUNTIME_GROUP TIDB_RUNTIME_HOME WORKER_LOG TIUP_DEPLOY_DIR TIKV_LOG
[ "$TIUP_CLUSTER_DEPLOYED" = "1" ] || die "state does not record a deployed TiUP cluster"

set -- $TIDB_HOSTS
TIDB_HOST="$1"
require_safe_name TIDB_RUNTIME_USER
require_safe_name TIDB_RUNTIME_GROUP
require_safe_name KEYSPACE_NAME
for name in TIDB_RUNTIME_HOME TIUP_DEPLOY_DIR WORKER_LOG TIKV_LOG; do
  require_safe_absolute_path "$name"
done
require_port TIDB_PORT
require_port TIKV_PORT
require_port TIDB_STATUS_PORT
require_port TIKV_STATUS_PORT

S3_ENDPOINT="https://s3.$REGION.amazonaws.com"

if [ "$apply" -eq 0 ]; then
  cat <<EOF
PLAN ONLY: validation commands will not be executed.
To validate, rerun: STATE_FILE="$STATE_FILE" $0 --apply

Will validate:
  cluster: $CLUSTER_NAME
  TiDB: $TIDB_HOST:$TIDB_PORT
  TiKV hosts: $TIKV_HOSTS
  TiKV instance IDs: $TIKV_IDS
  S3: s3://$S3_BUCKET/$S3_PREFIX/
  IAM profile: $TIKV_IAM_ROLE
  keyspace: $KEYSPACE_NAME
  topology: $TOPOLOGY_FILE
  worker config: $WORKER_CONFIG
EOF
  exit 0
fi

need_cmd aws
need_cmd tiup
need_cmd mysql
need_cmd curl
need_cmd grep
[ -f "$TOPOLOGY_FILE" ] || die "missing topology: $TOPOLOGY_FILE"
[ -f "$WORKER_CONFIG" ] || die "missing worker config: $WORKER_CONFIG"

if grep -Eq '<[^>]+>' "$TOPOLOGY_FILE" "$WORKER_CONFIG"; then
  die "topology or worker config still contains placeholders"
fi
if grep -Eq 's3-(key-id|secret-key)' "$TOPOLOGY_FILE" "$WORKER_CONFIG"; then
  die "canonical real-S3 config must omit static credentials"
fi

for expected in "$S3_ENDPOINT" "$S3_BUCKET" "$REGION" "$S3_PREFIX"; do
  grep -Fq "$expected" "$TOPOLOGY_FILE" || die "topology missing DFS value: $expected"
  grep -Fq "$expected" "$WORKER_CONFIG" || die "worker config missing DFS value: $expected"
done
grep -Fq "keyspace-name: \"$KEYSPACE_NAME\"" "$TOPOLOGY_FILE" || die "topology has the wrong TiDB keyspace"
grep -Fq "keyspace.pre-alloc: [\"$KEYSPACE_NAME\"]" "$TOPOLOGY_FILE" || die "topology has the wrong PD preallocated keyspace"

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"
public_block="$(aws --profile "$PROFILE" --region "$REGION" s3api get-public-access-block \
  --bucket "$S3_BUCKET" \
  --query 'PublicAccessBlockConfiguration.[BlockPublicAcls,IgnorePublicAcls,BlockPublicPolicy,RestrictPublicBuckets]' \
  --output text)"
[ "$public_block" = $'True\tTrue\tTrue\tTrue' ] || die "S3 public-access block is incomplete"
encryption="$(aws --profile "$PROFILE" --region "$REGION" s3api get-bucket-encryption \
  --bucket "$S3_BUCKET" \
  --query 'ServerSideEncryptionConfiguration.Rules[0].ApplyServerSideEncryptionByDefault.SSEAlgorithm' \
  --output text)"
[ "$encryption" = "AES256" ] || die "S3 bucket is not using the expected SSE-S3 encryption"
policy_document="$(aws --profile "$PROFILE" iam get-role-policy \
  --role-name "$TIKV_IAM_ROLE" --policy-name "$TIKV_IAM_POLICY" \
  --query PolicyDocument --output json)"
for expected_policy_value in "arn:aws:s3:::$S3_BUCKET" "arn:aws:s3:::$S3_BUCKET/$S3_PREFIX/*" \
  s3:ListBucket s3:PutObject s3:GetObject s3:DeleteObject; do
  printf '%s\n' "$policy_document" | grep -Fq "$expected_policy_value" || \
    die "TiKV IAM policy is missing $expected_policy_value"
done

word_count() {
  local values="$1"
  set -- $values
  printf '%s\n' "$#"
}

[ "$(word_count "$PD_IDS")" -eq "$(word_count "$PD_HOSTS")" ] || \
  die "PD instance and host inventories differ"
[ "$(word_count "$TIDB_IDS")" -eq "$(word_count "$TIDB_HOSTS")" ] || \
  die "TiDB instance and host inventories differ"
[ "$(word_count "$TIKV_IDS")" -eq "$(word_count "$TIKV_HOSTS")" ] || \
  die "TiKV instance and host inventories differ"

live_pd_hosts="$(aws --profile "$PROFILE" --region "$REGION" ec2 describe-instances \
  --instance-ids $PD_IDS \
  --query 'Reservations[].Instances[].PrivateIpAddress' --output text)"
live_tidb_hosts="$(aws --profile "$PROFILE" --region "$REGION" ec2 describe-instances \
  --instance-ids $TIDB_IDS \
  --query 'Reservations[].Instances[].PrivateIpAddress' --output text)"
live_tikv_hosts="$(aws --profile "$PROFILE" --region "$REGION" ec2 describe-instances \
  --instance-ids $TIKV_IDS \
  --query 'Reservations[].Instances[].PrivateIpAddress' --output text)"
for host in $PD_HOSTS; do
  case " $live_pd_hosts " in *" $host "*) ;; *) die "state PD host is not backed by a recorded EC2 ID: $host" ;; esac
done
for host in $TIDB_HOSTS; do
  case " $live_tidb_hosts " in *" $host "*) ;; *) die "state TiDB host is not backed by a recorded EC2 ID: $host" ;; esac
done
for host in $TIKV_HOSTS; do
  case " $live_tikv_hosts " in *" $host "*) ;; *) die "state TiKV host is not backed by a recorded EC2 ID: $host" ;; esac
done

display_output="$(run_bounded 120 tiup cluster display "$CLUSTER_NAME")"
printf '%s\n' "$display_output"
if printf '%s\n' "$display_output" | grep -Eq '[[:space:]](Down|Offline|N/A)[[:space:]]'; then
  die "one or more TiUP components are not Up"
fi
for host in $PD_HOSTS $TIDB_HOSTS $TIKV_HOSTS; do
  printf '%s\n' "$display_output" | grep -Fq "$host" || \
    die "recorded cluster host is absent from TiUP display: $host"
done

run_step "validate worker service state" 600 tiup cluster exec "$CLUSTER_NAME" -R tikv \
  --command 'systemctl is-active cse-worker'
run_step "validate worker runtime identity" 600 tiup cluster exec "$CLUSTER_NAME" -R tikv \
  --command "set -e
             test \"\$(systemctl show cse-worker -p User --value)\" = '$TIDB_RUNTIME_USER'
             test \"\$(systemctl show cse-worker -p Group --value)\" = '$TIDB_RUNTIME_GROUP'"
run_step "validate worker health" 600 tiup cluster exec "$CLUSTER_NAME" -R tikv \
  --command 'curl -sf http://127.0.0.1:9500/healthz'
run_step "validate TiKV S3 reachability" 600 tiup cluster exec "$CLUSTER_NAME" -R tikv \
  --command "code=\$(curl -sS --connect-timeout 10 -o /dev/null -w '%{http_code}' '$S3_ENDPOINT'); test \"\$code\" != 000"

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
expected_profile_count="$(printf '%s\n' $TIKV_IDS | wc -l | tr -d ' ')"
[ "$profile_count" -eq "$expected_profile_count" ] || die "not every TiKV has the expected profile"
expected_host_count="$(printf '%s\n' $TIKV_HOSTS | wc -l | tr -d ' ')"
[ "$expected_host_count" -eq "$expected_profile_count" ] || die "TiKV host and instance inventories differ"

mysql_cmd=(mysql -h "$TIDB_HOST" -P "$TIDB_PORT" -u "$MYSQL_USER")
if [ -n "$MYSQL_PASSWORD" ]; then
  export MYSQL_PWD="$MYSQL_PASSWORD"
fi

run_sql() {
  local sql="$1"
  run_bounded 120 "${mysql_cmd[@]}" -e "$sql"
}

validation_db="tidbx_validation_$(printf '%s' "$PROVISION_TOKEN" | tr '.-' '__')"
require_safe_name validation_db
cleanup_validation_db() {
  run_bounded 120 "${mysql_cmd[@]}" -e \
    "DROP DATABASE IF EXISTS \`$validation_db\`;" >/dev/null 2>&1 || true
}
trap cleanup_validation_db EXIT
trap 'exit 130' INT
trap 'exit 143' TERM
run_sql "DROP DATABASE IF EXISTS \`$validation_db\`;"

s3_object_count() {
  aws --profile "$PROFILE" --region "$REGION" s3api list-objects-v2 \
    --bucket "$S3_BUCKET" --prefix "$S3_PREFIX/" --max-keys 1 \
    --query KeyCount --output text
}

before_s3_count="$(s3_object_count)"
case "$before_s3_count" in ""|*[!0-9]*) die "invalid S3 object count: $before_s3_count" ;; esac

tidb_identity="$(run_sql "SELECT tidb_version();")"
printf '%s\n' "$tidb_identity"
printf '%s\n' "$tidb_identity" | grep -Fq 'Next Generation' || die "TiDB is not Next Generation"
resource_control="$(run_sql "SHOW VARIABLES LIKE 'tidb_enable_resource_control';")"
printf '%s\n' "$resource_control"
printf '%s\n' "$resource_control" | grep -Eq 'tidb_enable_resource_control[[:space:]]+ON' || die "resource control is not enabled"
fair_locking="$(run_sql "SET GLOBAL tidb_pessimistic_txn_fair_locking = OFF; SELECT @@global.tidb_pessimistic_txn_fair_locking;")"
printf '%s\n' "$fair_locking"
printf '%s\n' "$fair_locking" | tail -n 1 | grep -qx '0' || die "fair locking compatibility setting is not OFF"
run_sql "CREATE DATABASE \`$validation_db\`; CREATE TABLE \`$validation_db\`.smoke(a INT PRIMARY KEY); INSERT INTO \`$validation_db\`.smoke VALUES (1); SELECT * FROM \`$validation_db\`.smoke;"
analyze_output="$(run_sql "CREATE TABLE \`$validation_db\`.analyze_smoke (id INT PRIMARY KEY, v VARCHAR(32)); INSERT INTO \`$validation_db\`.analyze_smoke VALUES (1, 'ok'); ANALYZE TABLE \`$validation_db\`.analyze_smoke; SHOW STATS_META WHERE db_name='$validation_db' AND table_name='analyze_smoke';")"
printf '%s\n' "$analyze_output"
printf '%s\n' "$analyze_output" | grep -Fq 'analyze_smoke' || die "ANALYZE did not produce stats metadata"

for host in $TIKV_HOSTS; do
  remote_config="$(run_sql "SHOW CONFIG WHERE type='tikv' AND instance='$host:$TIKV_PORT' AND name LIKE 'kvengine.remote-%addr';")"
  for config_name in kvengine.remote-worker-addr kvengine.remote-coprocessor-addr; do
    printf '%s\n' "$remote_config" | grep -F "$config_name" | \
      grep -Fq 'http://127.0.0.1:9500/coprocessor' || \
      die "$host does not expose the expected $config_name"
  done
done

run_step "validate active worker S3 configuration" 600 tiup cluster exec "$CLUSTER_NAME" -R tikv --command \
  "set -e
   config='$TIDB_RUNTIME_HOME/cse-worker.toml'
   normalized=\$(sed 's/[[:space:]]//g' \"\$config\")
   printf '%s\n' \"\$normalized\" | grep -Fq 's3-endpoint=\"$S3_ENDPOINT\"'
   printf '%s\n' \"\$normalized\" | grep -Fq 's3-bucket=\"$S3_BUCKET\"'
   printf '%s\n' \"\$normalized\" | grep -Fq 's3-region=\"$REGION\"'
   printf '%s\n' \"\$normalized\" | grep -Fq 'prefix=\"$S3_PREFIX\"'"

run_step "validate PD identity and keyspace" 600 tiup cluster exec "$CLUSTER_NAME" -R pd --command \
  "binary=\$(find '$TIUP_DEPLOY_DIR' -type f -path '*/bin/pd-server' -print -quit)
   config=\$(find '$TIUP_DEPLOY_DIR' -type f -path '*/conf/pd.toml' -print -quit)
   test -n \"\$binary\" && test -n \"\$config\"
   \"\$binary\" -V 2>&1 | grep -F 'Kernel Type: Next Generation'
   grep -Fq '$KEYSPACE_NAME' \"\$config\""
run_step "validate TiDB identity and keyspace" 600 tiup cluster exec "$CLUSTER_NAME" -R tidb --command \
  "binary=\$(find '$TIUP_DEPLOY_DIR' -type f -path '*/bin/tidb-server' -print -quit)
   config=\$(find '$TIUP_DEPLOY_DIR' -type f -path '*/conf/tidb.toml' -print -quit)
   test -n \"\$binary\" && test -n \"\$config\"
   \"\$binary\" -V 2>&1 | grep -F 'Kernel Type: Next Generation'
   grep -Fq '$KEYSPACE_NAME' \"\$config\""
run_step "validate TiKV and worker build identity" 600 tiup cluster exec "$CLUSTER_NAME" -R tikv --command \
  "set -e
   tikv_binary=\$(find '$TIUP_DEPLOY_DIR' -type f -path '*/bin/tikv-server' -print -quit)
   worker_binary='$TIDB_RUNTIME_HOME/bin/tikv-worker'
   test -n \"\$tikv_binary\" && test -x \"\$worker_binary\"
   tikv_version=\$(\"\$tikv_binary\" -V 2>&1)
   worker_version=\$(\"\$worker_binary\" -V 2>&1)
   printf '%s\n' \"\$tikv_version\" | grep -F 'Cloud Storage Engine'
   printf '%s\n' \"\$worker_version\" | grep -F 'Cloud Storage Engine'
   tikv_commit=\$(printf '%s\n' \"\$tikv_version\" | sed -n 's/^Git Commit Hash:[[:space:]]*//p' | head -n 1)
   worker_commit=\$(printf '%s\n' \"\$worker_version\" | sed -n 's/^Git Commit Hash:[[:space:]]*//p' | head -n 1)
   test -n \"\$tikv_commit\" && test \"\$tikv_commit\" = \"\$worker_commit\""

error_pattern='AccessDenied|NoCredential|SignatureDoesNotMatch|invalid URL|remote compaction failed|Connection refused'
run_step "validate worker logs" 600 tiup cluster exec "$CLUSTER_NAME" -R tikv --command \
  "test -f '$WORKER_LOG' || exit 1
   if grep -Ei '$error_pattern' '$WORKER_LOG' | tail -n 20 | grep .; then exit 1; fi"
run_step "validate TiKV logs" 600 tiup cluster exec "$CLUSTER_NAME" -R tikv --command \
  "test -f '$TIKV_LOG' || exit 1
   if grep -Ei '$error_pattern|GetAllKeyspacesGCStates.*UNIMPLEMENTED' '$TIKV_LOG' | tail -n 20 | grep .; then exit 1; fi"

after_s3_count="$before_s3_count"
if [ "$before_s3_count" -eq 0 ]; then
  for _ in $(seq 1 60); do
    after_s3_count="$(s3_object_count)"
    case "$after_s3_count" in ""|*[!0-9]*) die "invalid S3 object count: $after_s3_count" ;; esac
    [ "$after_s3_count" -gt 0 ] && break
    sleep 5
  done
fi
[ "$after_s3_count" -gt 0 ] || \
  die "validation did not observe a cluster-written object in the fresh S3 prefix"

run_sql "DROP DATABASE \`$validation_db\`;"
remaining_schema="$(run_sql "SELECT SCHEMA_NAME FROM INFORMATION_SCHEMA.SCHEMATA WHERE SCHEMA_NAME='$validation_db';")"
printf '%s\n' "$remaining_schema" | grep -Fq "$validation_db" && \
  die "validation database still exists after cleanup"
trap - EXIT INT TERM
unset MYSQL_PWD

log "Canonical TiDB X validation completed"