#!/usr/bin/env bash
# pipeline/tools/local-e2e/bench/prepare.sh — prepare TPC-C data for all clusters in config.yaml
set -euo pipefail

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

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

if ! command -v tiup &>/dev/null; then
  echo "❌ tiup is not installed"
  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}"
WAREHOUSES="$(yaml_get "${cfg}" bench warehouses)"; WAREHOUSES="${WAREHOUSES:-2}"
THREADS="$(yaml_get "${cfg}" bench threads)";       THREADS="${THREADS:-4}"
DB_NAME="$(yaml_get "${cfg}" bench db)";            DB_NAME="${DB_NAME:-sbtest}"
TLS_MODE="$(yaml_get "${cfg}" bench tls)";          TLS_MODE="${TLS_MODE:-skip-verify}"
TIFLASH="$(yaml_get "${cfg}" bench tiflash)";       TIFLASH="${TIFLASH:-false}"

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

prepare_cluster() {
  local name="$1"
  local log="${LOG_DIR}/${name}-prepare.log"
  : > "${log}"

  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 tls;  tls="$(yaml_get  "${cfg}" clusters "${name}" tls)";  tls="${tls:-${TLS_MODE}}"
  local tiflash_cluster; tiflash_cluster="$(yaml_get "${cfg}" clusters "${name}" tiflash)"; tiflash_cluster="${tiflash_cluster:-${TIFLASH}}"

  echo "  [${name}] ${user}@${host}:${port}/${DB_NAME}  warehouses=${WAREHOUSES}"

  mysql --connect-timeout=15 \
        -h "${host}" -P "${port}" \
        -u "${user}" -p"${pass}" \
        --ssl-mode=REQUIRED \
        -e "CREATE DATABASE IF NOT EXISTS \`${DB_NAME}\`;" 2>>"${log}" \
    || { echo "  [${name}] ⚠️  failed to connect, skipping"; return 1; }

  echo "  [${name}] prepare..."
  tiup bench tpcc \
    --host "${host}" --port "${port}" \
    --user "${user}" --password "${pass}" \
    --db "${DB_NAME}" --conn-params "tls=${tls}" \
    --warehouses "${WAREHOUSES}" --threads "${THREADS}" \
    --ignore-error \
    prepare >>"${log}" 2>&1 \
    && echo "  [${name}] ✅ prepare done  →  ${log}" \
    || echo "  [${name}] ⚠️  prepare failed (see ${log})"

  if [[ "${tiflash_cluster}" == "true" ]]; then
    echo "  [${name}] adding TiFlash replicas..."
    local mysql_opts=( --connect-timeout=15 -h "${host}" -P "${port}" -u "${user}" -p"${pass}" --ssl-mode=REQUIRED )
    for tbl in "${TPCC_TABLES[@]}"; do
      mysql "${mysql_opts[@]}" \
        -e "ALTER TABLE \`${DB_NAME}\`.\`${tbl}\` SET TIFLASH REPLICA 1;" >>"${log}" 2>&1 \
        && echo "  [${name}]   ✅ ${tbl}" \
        || echo "  [${name}]   ⚠️  ${tbl} failed (see ${log})"
    done
    echo "  [${name}] ⏳ waiting for TiFlash sync (up to 5 minutes)..."
    local deadline=$(( $(date +%s) + 300 ))
    local all_synced=false
    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
        all_synced=true
        break
      fi
      sleep 10
    done
    ${all_synced} \
      && echo "  [${name}] ✅ TiFlash replica sync completed" \
      || echo "  [${name}] ⚠️  TiFlash replica sync timed out; please check manually later"
  fi
}

echo "================================================================"
echo " TPC-C prepare  config: $(basename "${CONFIG_FILE}")"
echo "================================================================"
echo "  warehouses=${WAREHOUSES}  threads=${THREADS}"
echo ""

declare -a PIDS=()
while IFS= read -r name; do
  [[ -z "${name}" ]] && continue
  prepare_cluster "${name}" &
  PIDS+=($!)
done < <(yaml_cluster_names "${cfg}")

ALL_OK=true
for pid in "${PIDS[@]}"; do
  wait "${pid}" || ALL_OK=false
done

echo ""
${ALL_OK} && echo "✅ Prepare completed for all clusters" || { echo "⚠️  Some clusters failed"; exit 1; }