#!/usr/bin/env bash
# bench/run.sh - read config.yaml and run TPC-C load against TiDB Cloud clusters
#
# Usage:
#   bash pipeline/tools/local-e2e/bench/run.sh
#   BENCH_CONFIG=my.yaml bash pipeline/tools/local-e2e/bench/run.sh
#   BENCH_SKIP_PREPARE=true bash pipeline/tools/local-e2e/bench/run.sh
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}"
  echo "   First run: cp pipeline/tools/local-e2e/bench/config.yaml.template pipeline/tools/local-e2e/bench/config.yaml"
  exit 1
fi

if ! command -v tiup &>/dev/null; then
  echo "❌ tiup is not installed"
  exit 1
fi

# Read nested fields from YAML via python3
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
}

# Read global bench settings
cfg="${CONFIG_FILE}"

WAREHOUSES="$(yaml_get "${cfg}" bench warehouses)"; WAREHOUSES="${WAREHOUSES:-2}"
THREADS="$(yaml_get "${cfg}" bench threads)";       THREADS="${THREADS:-4}"
TIME_LIMIT="$(yaml_get "${cfg}" bench time)";       TIME_LIMIT="${TIME_LIMIT:-5m}"
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}"
SKIP_PREPARE="$(yaml_get "${cfg}" bench skip_prepare)"; SKIP_PREPARE="${BENCH_SKIP_PREPARE:-${SKIP_PREPARE:-false}}"

# Per-cluster run function
run_bench_for_cluster() {
  local name="$1"
  local log="${LOG_DIR}/${name}.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)"
  # per-cluster tls override
  local tls;  tls="$(yaml_get  "${cfg}" clusters "${name}" tls)";  tls="${tls:-${TLS_MODE}}"

  echo "  [${name}] ${user}@${host}:${port}/${DB_NAME}  warehouses=${WAREHOUSES} threads=${THREADS} time=${TIME_LIMIT} tls=${tls}"

  # ── auto-create database ──────────────────────────────────────────────────
  echo "  [${name}] creating database ${DB_NAME} if not exists..."
  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 and create database, skipping"; return 1; }

  local base_args=(
    --host       "${host}"
    --port       "${port}"
    --user       "${user}"
    --password   "${pass}"
    --db         "${DB_NAME}"
    --conn-params "tls=${tls}"
    --warehouses "${WAREHOUSES}"
    --threads    "${THREADS}"
    --ignore-error
  )

  if [[ "${SKIP_PREPARE}" != "true" ]]; then
    echo "  [${name}] prepare..."
    tiup bench tpcc "${base_args[@]}" prepare >>"${log}" 2>&1 \
      || echo "  [${name}] ⚠️  prepare failed (tables may already exist, continuing to run)"
  fi

  echo "  [${name}] run (${TIME_LIMIT})..."
  tiup bench tpcc "${base_args[@]}" --time "${TIME_LIMIT}" run >>"${log}" 2>&1
  echo "  [${name}] ✅ done  →  ${log}"
}

# Main flow: run all clusters with configured hosts in parallel
echo "================================================================"
echo " TPC-C Load  config: $(basename "${CONFIG_FILE}")"
echo "================================================================"
echo "  warehouses=${WAREHOUSES}  threads=${THREADS}  time=${TIME_LIMIT}"
echo ""

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

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

# Print per-cluster result summary
echo ""
echo "================================================================"
echo " Summary"
echo "================================================================"
for name in "${NAMES[@]}"; do
  log="${LOG_DIR}/${name}.log"
  echo ""
  echo "── ${name} ──────────────────────────────────────────────────"
  if [[ -f "${log}" ]]; then
    # show tpmC / tpmTotal lines from the run output
    grep -E "(tpmC|tpmTotal|Finished|PASS|FAIL|ERROR)" "${log}" 2>/dev/null || tail -5 "${log}"
  fi
done

echo ""
${ALL_OK} && echo "✅ All runs completed" || { echo "⚠️  Some clusters failed"; exit 1; }