#!/usr/bin/env bash
# pipeline/tools/local-e2e/bench/tiflash.sh — enable TiFlash replicas on TPC-C tables for clusters in config.yaml
#
# Usage:
#   bash pipeline/tools/local-e2e/bench/tiflash.sh                  # only clusters with tiflash: true in config
#   bash pipeline/tools/local-e2e/bench/tiflash.sh --all            # all clusters regardless of flag
#   bash pipeline/tools/local-e2e/bench/tiflash.sh dedicated        # specific cluster by name
set -euo pipefail

BENCH_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
CONFIG_FILE="${BENCH_CONFIG:-${BENCH_DIR}/config.yaml}"

if [[ ! -f "${CONFIG_FILE}" ]]; then
  echo "❌ Config file not found: ${CONFIG_FILE}"
  exit 1
fi

yaml_get() {
  python3 - "$@" <<'PYEOF'
import sys, yaml
def get_nested(d, *keys):
    for k in keys:
        if not isinstance(d, dict) or k not in d:
            return ""
        d = d[k]
    return "" if d is None else str(d)
with open(sys.argv[1]) as f:
    data = yaml.safe_load(f)
print(get_nested(data, *sys.argv[2:]))
PYEOF
}

yaml_cluster_names() {
  python3 - "$1" <<'PYEOF'
import sys, yaml
with open(sys.argv[1]) as f:
    data = yaml.safe_load(f)
for k in (data.get("clusters") or {}).keys():
    print(k)
PYEOF
}

cfg="${CONFIG_FILE}"
DB_NAME="$(yaml_get "${cfg}" bench db)"; DB_NAME="${DB_NAME:-sbtest}"
GLOBAL_TIFLASH="$(yaml_get "${cfg}" bench tiflash)"; GLOBAL_TIFLASH="${GLOBAL_TIFLASH:-false}"

TPCC_TABLES=(warehouse district customer history new_order orders order_line stock item)

FORCE_ALL=false
TARGET_CLUSTER=""
for arg in "$@"; do
  case "${arg}" in
    --all) FORCE_ALL=true ;;
    *)     TARGET_CLUSTER="${arg}" ;;
  esac
done

enable_tiflash() {
  local name="$1"

  local host; host="$(yaml_get "${cfg}" clusters "${name}" host)"
  if [[ -z "${host}" ]]; then
    echo "  [${name}] host is not configured, skipping"
    return 0
  fi

  local port; port="$(yaml_get "${cfg}" clusters "${name}" port)"; port="${port:-4000}"
  local user; user="$(yaml_get "${cfg}" clusters "${name}" user)"; user="${user:-root}"
  local pass; pass="$(yaml_get "${cfg}" clusters "${name}" password)"

  local mysql_opts=( --connect-timeout=15 -h "${host}" -P "${port}" -u "${user}" -p"${pass}" --ssl-mode=REQUIRED )

  echo "  [${name}] ALTER TABLE ... SET TIFLASH REPLICA 1 on \`${DB_NAME}\`..."
  for tbl in "${TPCC_TABLES[@]}"; do
    if mysql "${mysql_opts[@]}" \
        -e "ALTER TABLE \`${DB_NAME}\`.\`${tbl}\` SET TIFLASH REPLICA 1;" 2>/dev/null; then
      echo "  [${name}]   ✅ ${tbl}"
    else
      echo "  [${name}]   ⚠️  ${tbl} failed (replica may already exist or the cluster may not support TiFlash)"
    fi
  done

  echo "  [${name}] ⏳ waiting for TiFlash sync (up to 5 minutes)..."
  local deadline=$(( $(date +%s) + 300 ))
  while [[ $(date +%s) -lt ${deadline} ]]; do
    local not_ready
    not_ready=$(mysql "${mysql_opts[@]}" -sNe \
      "SELECT COUNT(*) FROM information_schema.tiflash_replica
       WHERE TABLE_SCHEMA='${DB_NAME}' AND AVAILABLE=0;" 2>/dev/null || echo "999")
    if [[ "${not_ready}" == "0" ]]; then
      echo "  [${name}] ✅ TiFlash replica sync completed"
      return 0
    fi
    sleep 10
  done
  echo "  [${name}] ⚠️  TiFlash replica sync timed out; please check manually later"
}

echo "================================================================"
echo " TiFlash Replica Setup  db: ${DB_NAME}"
echo "================================================================"
echo ""

if [[ -n "${TARGET_CLUSTER}" ]]; then
  enable_tiflash "${TARGET_CLUSTER}"
else
  while IFS= read -r name; do
    [[ -z "${name}" ]] && continue
    local_tiflash="$(yaml_get "${cfg}" clusters "${name}" tiflash)"
    local_tiflash="${local_tiflash:-${GLOBAL_TIFLASH}}"
    if [[ "${FORCE_ALL}" == "true" || "${local_tiflash}" == "true" ]]; then
      enable_tiflash "${name}"
    else
      echo "  [${name}] tiflash=false, skipping (use --all to force)"
    fi
  done < <(yaml_cluster_names "${cfg}")
fi

echo ""
echo "✅ Done"