#!/bin/bash

log() {
  printf '[%s] %s\n' "$(date -u '+%Y-%m-%dT%H:%M:%SZ')" "$*" >&2
}

die() {
  printf 'ERROR: %s\n' "$*" >&2
  exit 1
}

need_cmd() {
  local cmd="$1"
  command -v "$cmd" >/dev/null 2>&1 || die "missing required command: $cmd"
}

require_env() {
  local missing=()
  local name
  for name in "$@"; do
    if [ -z "${!name:-}" ]; then
      missing+=("$name")
    fi
  done
  if [ "${#missing[@]}" -gt 0 ]; then
    die "missing required environment variables: ${missing[*]}"
  fi
}

require_positive_int() {
  local name="$1"
  local value="${!name:-}"
  case "$value" in
    ""|*[!0-9]*) die "$name must be a positive integer" ;;
  esac
  [ "$value" -gt 0 ] || die "$name must be greater than zero"
}

require_nonnegative_int() {
  local name="$1"
  local value="${!name:-}"
  case "$value" in
    ""|*[!0-9]*) die "$name must be a non-negative integer" ;;
  esac
}

require_nonnegative_decimal() {
  local name="$1"
  local value="${!name:-}"
  case "$value" in
    ""|*[!0-9.]*|.*|*.) die "$name must be a non-negative decimal" ;;
  esac
  case "$value" in
    *.*.*) die "$name must contain at most one decimal point" ;;
  esac
}

require_no_newline() {
  local name="$1"
  local value="${!name:-}"
  case "$value" in
    *$'\n'*|*$'\r'*) die "$name must not contain a newline" ;;
  esac
}

require_port() {
  local name="$1"
  require_positive_int "$name"
  [ "${!name}" -le 65535 ] || die "$name must be at most 65535"
}

require_safe_name() {
  local name="$1"
  local value="${!name:-}"
  case "$value" in
    ""|*[!A-Za-z0-9_.-]*) die "$name contains unsupported characters" ;;
  esac
}

require_safe_absolute_path() {
  local name="$1"
  local value="${!name:-}"
  case "$value" in
    /*) ;;
    *) die "$name must be an absolute path" ;;
  esac
  case "$value" in
    *[!A-Za-z0-9_./-]*) die "$name contains unsupported characters" ;;
  esac
}

timeout_cmd() {
  if command -v timeout >/dev/null 2>&1; then
    printf 'timeout\n'
    return 0
  fi
  if command -v gtimeout >/dev/null 2>&1; then
    printf 'gtimeout\n'
    return 0
  fi
  return 1
}

run_step_log_file() {
  local name="$1"
  local log_dir="${RUN_LOG_DIR:-./run-logs}"
  local safe_name
  local timestamp

  mkdir -p "$log_dir"
  safe_name="$(printf '%s' "$name" | tr -cs '[:alnum:]_.-' '_' | sed 's/^_//; s/_$//')"
  if [ -z "$safe_name" ]; then
    safe_name="step"
  fi
  timestamp="$(date -u '+%Y%m%dT%H%M%SZ')"
  printf '%s/%s_%s.log\n' "$log_dir" "$timestamp" "$safe_name"
}

run_step_diagnostics() {
  local name="$1"
  local status="$2"
  local log_file="$3"

  log "===== ${name} failed status=${status}; log=${log_file} ====="
  case "$status" in
    124|137|143)
      log "step likely timed out or was killed; watchdog kill-after=${RUN_STEP_KILL_AFTER:-30s}"
      ;;
  esac

  if [ -n "${RUN_STEP_DIAGNOSTIC_CMD:-}" ]; then
    log "===== ${name} diagnostics ====="
    bash -lc "$RUN_STEP_DIAGNOSTIC_CMD" 2>&1 | tee -a "$log_file" || true
  fi
}

run_step() {
  local name="$1"
  local timeout_seconds="$2"
  shift 2

  local cmd
  local log_file
  local status
  local had_errexit=0

  cmd="$(timeout_cmd)" || die "install GNU timeout/coreutils before running long steps"
  log_file="$(run_step_log_file "$name")"
  log "===== ${name} timeout=${timeout_seconds}s log=${log_file} ====="

  case "$-" in
    *e*)
      had_errexit=1
      set +e
      ;;
  esac

  "$cmd" --kill-after="${RUN_STEP_KILL_AFTER:-30s}" --preserve-status "$timeout_seconds" "$@" \
    2>&1 | tee -a "$log_file"
  status=${PIPESTATUS[0]}

  if [ "$had_errexit" -eq 1 ]; then
    set -e
  fi

  if [ "$status" -ne 0 ]; then
    run_step_diagnostics "$name" "$status" "$log_file"
    return "$status"
  fi

  log "===== ${name} completed log=${log_file} ====="
}

run_bounded() {
  local timeout_seconds="$1"
  shift
  local cmd
  cmd="$(timeout_cmd)" || die "install GNU timeout/coreutils before running bounded commands"
  "$cmd" --kill-after="${RUN_STEP_KILL_AFTER:-30s}" --preserve-status "$timeout_seconds" "$@"
}

write_state() {
  local state_file="$1"
  local key="$2"
  local value="$3"
  local tmp_file
  local encoded
  local previous_umask

  need_cmd base64
  case "$key" in
    ""|[!A-Z_]*|*[!A-Z0-9_]*) die "invalid state key: $key" ;;
  esac
  mkdir -p "$(dirname "$state_file")"
  previous_umask="$(umask)"
  umask 077
  tmp_file="$(mktemp "${state_file}.tmp.XXXXXX")"
  if [ -f "$state_file" ]; then
    grep -v "^${key}=" "$state_file" > "$tmp_file" || true
  fi
  encoded="$(printf '%s' "$value" | base64 | tr -d '\n')"
  printf '%s=base64:%s\n' "$key" "$encoded" >> "$tmp_file"
  mv "$tmp_file" "$state_file"
  umask "$previous_umask"
}

load_state_file() {
  local state_file="$1"
  local key
  local encoded
  local value
  local seen_keys=""
  [ -f "$state_file" ] || die "missing required state file: $state_file"
  [ ! -L "$state_file" ] || die "state file must not be a symbolic link: $state_file"
  need_cmd base64
  while IFS='=' read -r key encoded; do
    [ -n "$key" ] || continue
    case "$key" in
      [!A-Z_]*|*[!A-Z0-9_]*) die "invalid key in state file: $key" ;;
    esac
    case "$key" in
      ALL_CLIENT_TOKENS|ALL_INSTANCE_IDS|ALL_VOLUME_IDS|CLEANUP_OWNER|CLUSTER_KNOWN_HOSTS|CLUSTER_NAME|CLUSTER_PAUSED|CLUSTER_SG|CLUSTER_SG_NAME|CLUSTER_SSH_KEY|CONTROL_HOST_ID|CONTROL_SG|CREATED_BY|DEPLOY_SSH_KEY|DEPLOY_SSH_USER|EXPECTED_AWS_ACCOUNT_ID|GRAFANA_PORT|KEYSPACE_NAME|MONITOR_HOSTS|MONITOR_IDS|PAUSE_TIUP_STOPPED|PAUSE_WORKERS_STOPPED|PD_CLIENT_PORT|PD_HOSTS|PD_IDS|PD_PEER_PORT|PD_TARBALL|PRESERVED_S3_BUCKET|PRESSURE_HOSTS|PRESSURE_IDS|PROFILE|PROMETHEUS_PORT|PROVISION_TOKEN|REGION|RUN_DIR|RUN_LOG_DIR|S3_BUCKET|S3_BUCKET_CREATED|S3_NETWORK_MODE|S3_NETWORK_POLICY_APPROVED|S3_NETWORK_RESOURCE_ID|S3_NETWORK_ROUTE_TABLE_ID|S3_PREFIX|SG|SSH_PORT|SUBNET_ID|TIDB_HOSTS|TIDB_IDS|TIDB_PORT|TIDB_RUNTIME_GROUP|TIDB_RUNTIME_HOME|TIDB_RUNTIME_USER|TIDB_STATUS_PORT|TIDB_TARBALL|TIKV_HOSTS|TIKV_IAM_POLICY|TIKV_IAM_PROFILE|TIKV_IAM_PROFILE_ROLE_ATTACHED|TIKV_IAM_ROLE|TIKV_IDS|TIKV_PORT|TIKV_STATUS_PORT|TIKV_TARBALL|TIUP_CLUSTER_DEPLOYED|TIUP_CLUSTER_DEPLOY_INTENT|TIUP_DATA_DIR|TIUP_DEPLOY_DIR|TIUP_METADATA_VERSION|TOPOLOGY_FILE|VPC_ID|WORKER_CONFIG|WORKER_TARBALL) ;;
      *) die "unsupported key in state file: $key" ;;
    esac
    case " $seen_keys " in *" $key "*) die "duplicate key in state file: $key" ;; esac
    seen_keys="${seen_keys:+$seen_keys }$key"
    case "$encoded" in
      base64:*) ;;
      *) die "invalid encoding for $key in state file" ;;
    esac
    value="$(printf '%s' "${encoded#base64:}" | base64 --decode)" || \
      die "invalid base64 value for $key in state file"
    printf -v "$key" '%s' "$value"
  done < "$state_file"
}

unset_vars() {
  local name
  for name in "$@"; do
    unset "$name"
  done
}

sha256_stream() {
  need_cmd sha256sum
  sha256sum | awk '{print $1}'
}

aws_approval_manifest() {
  local name
  for name in \
    PROFILE EXPECTED_AWS_ACCOUNT_ID REGION AZ VPC_ID SUBNET_ID AMI KEY_NAME \
    RUN_DIR STATE_FILE RUN_LOG_DIR CLUSTER_NAME PROVISION_TOKEN OWNER CREATED_BY \
    CLEANUP_OWNER EXPIRES_AT COST_CEILING_USD ESTIMATED_BURN_RATE_USD_PER_HOUR \
    ESTIMATED_BURN_RATE_USD_PER_DAY ESTIMATED_TOTAL_COST_USD CONTROL_HOST_ID \
    CONTROL_TYPE CONTROL_SG 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 PRESSURE_COUNT PRESSURE_TYPE PRESSURE_ROOT_GB S3_BUCKET \
    S3_PREFIX TIKV_IAM_ROLE TIKV_IAM_POLICY OPERATOR_CIDR ASSOCIATE_PUBLIC_IP \
    SSH_PORT TIDB_PORT GRAFANA_PORT PROMETHEUS_PORT S3_NETWORK_MODE \
    S3_NETWORK_RESOURCE_ID S3_NETWORK_POLICY_APPROVED NETWORK_EXTRAS; do
    require_no_newline "$name"
    printf '%s=%s\n' "$name" "${!name:-}"
  done
}

aws_approval_sha256() {
  aws_approval_manifest | sha256_stream
}

validate_s3_network_selection() {
  local mode="$1"
  local public_ip="$2"
  local resource_id="$3"

  case "$public_ip" in
    0|1) ;;
    *) die "ASSOCIATE_PUBLIC_IP must be 0 or 1" ;;
  esac

  case "$mode" in
    gateway-endpoint)
      [ "$public_ip" = "0" ] || die "gateway-endpoint mode requires ASSOCIATE_PUBLIC_IP=0"
      case "$resource_id" in vpce-*) ;; *) die "gateway-endpoint mode requires an S3 VPC endpoint ID" ;; esac
      ;;
    nat)
      [ "$public_ip" = "0" ] || die "nat mode requires ASSOCIATE_PUBLIC_IP=0"
      case "$resource_id" in nat-*) ;; *) die "nat mode requires a NAT gateway ID" ;; esac
      ;;
    public-ip)
      [ "$public_ip" = "1" ] || die "public-ip mode requires ASSOCIATE_PUBLIC_IP=1"
      case "$resource_id" in igw-*) ;; *) die "public-ip mode requires an internet gateway ID" ;; esac
      ;;
    *) die "S3_NETWORK_MODE must be gateway-endpoint, nat, or public-ip" ;;
  esac
}

resolve_subnet_route_table() {
  local profile="$1"
  local region="$2"
  local vpc_id="$3"
  local subnet_id="$4"
  local route_table_id

  route_table_id="$(aws --profile "$profile" --region "$region" ec2 describe-route-tables \
    --filters "Name=association.subnet-id,Values=$subnet_id" \
    --query 'RouteTables[0].RouteTableId' --output text)"
  if [ -z "$route_table_id" ] || [ "$route_table_id" = "None" ]; then
    route_table_id="$(aws --profile "$profile" --region "$region" ec2 describe-route-tables \
      --filters "Name=vpc-id,Values=$vpc_id" "Name=association.main,Values=true" \
      --query 'RouteTables[0].RouteTableId' --output text)"
  fi
  [ -n "$route_table_id" ] && [ "$route_table_id" != "None" ] || \
    die "could not resolve a route table for subnet $subnet_id"
  printf '%s\n' "$route_table_id"
}

validate_subnet_az() {
  local profile="$1"
  local region="$2"
  local subnet_id="$3"
  local expected_az="$4"
  local actual_az

  actual_az="$(aws --profile "$profile" --region "$region" ec2 describe-subnets \
    --subnet-ids "$subnet_id" --query 'Subnets[0].AvailabilityZone' --output text)"
  [ "$actual_az" = "$expected_az" ] || \
    die "subnet $subnet_id is in $actual_az, not approved AZ $expected_az"
}

resolve_ami_root_device() {
  local profile="$1"
  local region="$2"
  local ami="$3"
  local image
  local architecture
  local root_device_name
  local root_device_type
  local state

  image="$(aws --profile "$profile" --region "$region" ec2 describe-images \
    --image-ids "$ami" \
    --query 'Images[0].[Architecture,RootDeviceName,RootDeviceType,State]' --output text)"
  read -r architecture root_device_name root_device_type state <<< "$image"
  [ "$architecture" = "x86_64" ] || die "AMI $ami architecture is $architecture, expected x86_64"
  [ "$root_device_type" = "ebs" ] || die "AMI $ami root device type is $root_device_type, expected ebs"
  [ "$state" = "available" ] || die "AMI $ami state is $state, expected available"
  case "$root_device_name" in /dev/*) ;; *) die "AMI $ami has invalid root device name: $root_device_name" ;; esac
  printf '%s\n' "$root_device_name"
}

validate_s3_network_path() {
  local profile="$1"
  local region="$2"
  local vpc_id="$3"
  local subnet_id="$4"
  local mode="$5"
  local public_ip="$6"
  local resource_id="$7"
  local route_table_id
  local actual
  local state
  local service
  local route_tables
  local nat_subnet_id
  local nat_route_table_id
  local internet_gateway_id

  validate_s3_network_selection "$mode" "$public_ip" "$resource_id"
  route_table_id="$(resolve_subnet_route_table "$profile" "$region" "$vpc_id" "$subnet_id")"

  case "$mode" in
    gateway-endpoint)
      actual="$(aws --profile "$profile" --region "$region" ec2 describe-vpc-endpoints \
        --vpc-endpoint-ids "$resource_id" --query 'VpcEndpoints[0].VpcId' --output text)"
      state="$(aws --profile "$profile" --region "$region" ec2 describe-vpc-endpoints \
        --vpc-endpoint-ids "$resource_id" --query 'VpcEndpoints[0].State' --output text)"
      service="$(aws --profile "$profile" --region "$region" ec2 describe-vpc-endpoints \
        --vpc-endpoint-ids "$resource_id" --query 'VpcEndpoints[0].ServiceName' --output text)"
      route_tables="$(aws --profile "$profile" --region "$region" ec2 describe-vpc-endpoints \
        --vpc-endpoint-ids "$resource_id" --query 'VpcEndpoints[0].RouteTableIds' --output text)"
      [ "$actual" = "$vpc_id" ] || die "S3 VPC endpoint is not in VPC $vpc_id"
      [ "$state" = "available" ] || die "S3 VPC endpoint is not available"
      [ "$service" = "com.amazonaws.$region.s3" ] || die "VPC endpoint is not the regional S3 service"
      case " $route_tables " in
        *" $route_table_id "*) ;;
        *) die "S3 VPC endpoint is not associated with route table $route_table_id" ;;
      esac
      ;;
    nat)
      actual="$(aws --profile "$profile" --region "$region" ec2 describe-nat-gateways \
        --nat-gateway-ids "$resource_id" --query 'NatGateways[0].VpcId' --output text)"
      state="$(aws --profile "$profile" --region "$region" ec2 describe-nat-gateways \
        --nat-gateway-ids "$resource_id" --query 'NatGateways[0].State' --output text)"
      [ "$actual" = "$vpc_id" ] || die "NAT gateway is not in VPC $vpc_id"
      [ "$state" = "available" ] || die "NAT gateway is not available"
      actual="$(aws --profile "$profile" --region "$region" ec2 describe-route-tables \
        --route-table-ids "$route_table_id" \
        --query "RouteTables[0].Routes[?DestinationCidrBlock=='0.0.0.0/0'].NatGatewayId | [0]" \
        --output text)"
      [ "$actual" = "$resource_id" ] || die "route table $route_table_id does not use NAT gateway $resource_id"
      nat_subnet_id="$(aws --profile "$profile" --region "$region" ec2 describe-nat-gateways \
        --nat-gateway-ids "$resource_id" --query 'NatGateways[0].SubnetId' --output text)"
      nat_route_table_id="$(resolve_subnet_route_table "$profile" "$region" "$vpc_id" "$nat_subnet_id")"
      internet_gateway_id="$(aws --profile "$profile" --region "$region" ec2 describe-route-tables \
        --route-table-ids "$nat_route_table_id" \
        --query "RouteTables[0].Routes[?DestinationCidrBlock=='0.0.0.0/0'].GatewayId | [0]" \
        --output text)"
      case "$internet_gateway_id" in
        igw-*) ;;
        *) die "NAT subnet route table $nat_route_table_id has no internet gateway default route" ;;
      esac
      state="$(aws --profile "$profile" --region "$region" ec2 describe-internet-gateways \
        --internet-gateway-ids "$internet_gateway_id" \
        --query "InternetGateways[0].Attachments[?VpcId=='$vpc_id'].State | [0]" --output text)"
      [ "$state" = "available" ] || die "NAT internet gateway is not attached to VPC $vpc_id"
      ;;
    public-ip)
      state="$(aws --profile "$profile" --region "$region" ec2 describe-internet-gateways \
        --internet-gateway-ids "$resource_id" \
        --query "InternetGateways[0].Attachments[?VpcId=='$vpc_id'].State | [0]" --output text)"
      [ "$state" = "available" ] || die "internet gateway is not attached to VPC $vpc_id"
      actual="$(aws --profile "$profile" --region "$region" ec2 describe-route-tables \
        --route-table-ids "$route_table_id" \
        --query "RouteTables[0].Routes[?DestinationCidrBlock=='0.0.0.0/0'].GatewayId | [0]" \
        --output text)"
      [ "$actual" = "$resource_id" ] || die "route table $route_table_id does not use internet gateway $resource_id"
      ;;
  esac

  printf '%s\n' "$route_table_id"
}

export -f log
export -f die
export -f need_cmd
export -f require_env
export -f require_positive_int
export -f require_nonnegative_int
export -f require_nonnegative_decimal
export -f require_port
export -f require_safe_name
export -f require_safe_absolute_path
export -f require_no_newline
export -f timeout_cmd
export -f run_step_log_file
export -f run_step_diagnostics
export -f run_step
export -f run_bounded
export -f write_state
export -f load_state_file
export -f unset_vars
export -f sha256_stream
export -f aws_approval_manifest
export -f aws_approval_sha256
export -f validate_s3_network_selection
export -f resolve_subnet_route_table
export -f validate_subnet_az
export -f resolve_ami_root_device
export -f validate_s3_network_path