#!/usr/bin/env bash
# urus-bench/run_all.sh
#
# Runs the full (server, scenario) matrix and writes a summary.md.
#
# Default matrix (8 runs):
#   urus-mem    × s1, s2, s3
#   axum-mem    × s1, s2, s3
#   urus-sqlite × s4
#   axum-sqlite × s4
#
# Cowboy is not built; skipped.
#
# Usage:
#   ./run_all.sh                          # full matrix, defaults
#   ./run_all.sh --quick                  # short timings (15s+15s+15s)
#   ./run_all.sh --scenarios=s1,s2        # subset of scenarios
#   ./run_all.sh --servers=urus-mem       # subset of servers
#   ./run_all.sh --batch=foo              # batch label (default: timestamp)
#
# Per-run env vars (PROBE_SEC, WARMUP_SEC, MEASURE_SEC, PREPOP, SERVER_CPUS,
# LOADGEN_CPUS, WRK_CONNS, WRK_THREADS, SAT_RATIO) are passed through to
# runner.sh; --quick is just a preset for the timings.
#
# Output:
#   results/<batch>/                 batch dir
#     <server>-<scenario>/           per-run dir (from runner.sh)
#     summary.md                     aggregated headline table
#     index.json                     machine-readable index of all runs

set -euo pipefail

REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

# ---------------------------------------------------------------------------
# Defaults
# ---------------------------------------------------------------------------

ALL_SERVERS_MEM=(urus-mem axum-mem)
ALL_SERVERS_SQLITE=(urus-sqlite axum-sqlite)
ALL_SCENARIOS_MEM=(s1 s2 s3)
ALL_SCENARIOS_SQLITE=(s4)

SERVERS_FILTER=""
SCENARIOS_FILTER=""
BATCH=""
QUICK=0

# ---------------------------------------------------------------------------
# Arg parsing
# ---------------------------------------------------------------------------

for arg in "$@"; do
  case "$arg" in
    --quick)            QUICK=1 ;;
    --servers=*)        SERVERS_FILTER="${arg#*=}" ;;
    --scenarios=*)      SCENARIOS_FILTER="${arg#*=}" ;;
    --batch=*)          BATCH="${arg#*=}" ;;
    -h|--help)
      sed -n '2,/^$/p' "$0" | sed 's/^# \{0,1\}//'
      exit 0 ;;
    *)
      echo "run_all: unknown arg: $arg" >&2
      exit 2 ;;
  esac
done

if [[ "$QUICK" == "1" ]]; then
  export PROBE_SEC="${PROBE_SEC:-15}"
  export WARMUP_SEC="${WARMUP_SEC:-15}"
  export MEASURE_SEC="${MEASURE_SEC:-15}"
fi

BATCH="${BATCH:-$(date +%Y%m%d-%H%M%S)}"
BATCH_DIR="${REPO_ROOT}/results/${BATCH}"
mkdir -p "$BATCH_DIR"

# ---------------------------------------------------------------------------
# Build the matrix
# ---------------------------------------------------------------------------
#
# A run is a pair "<server> <scenario>". We compute the full matrix, then
# apply filters. s1/s2/s3 always pair with the memory backends; s4 only
# pairs with sqlite backends.

contains() {
  # contains "needle" "csv-list" → 0 if needle in csv (or csv empty), else 1
  local needle="$1" haystack="$2"
  [[ -z "$haystack" ]] && return 0
  local IFS=,
  for item in $haystack; do
    [[ "$item" == "$needle" ]] && return 0
  done
  return 1
}

runs=()
for srv in "${ALL_SERVERS_MEM[@]}"; do
  contains "$srv" "$SERVERS_FILTER" || continue
  for sc in "${ALL_SCENARIOS_MEM[@]}"; do
    contains "$sc" "$SCENARIOS_FILTER" || continue
    runs+=("$srv $sc")
  done
done
for srv in "${ALL_SERVERS_SQLITE[@]}"; do
  contains "$srv" "$SERVERS_FILTER" || continue
  for sc in "${ALL_SCENARIOS_SQLITE[@]}"; do
    contains "$sc" "$SCENARIOS_FILTER" || continue
    runs+=("$srv $sc")
  done
done

if [[ "${#runs[@]}" -eq 0 ]]; then
  echo "run_all: empty matrix (filters too narrow?)" >&2
  exit 2
fi

# ---------------------------------------------------------------------------
# Build binaries up front
# ---------------------------------------------------------------------------
#
# We always rebuild in --release. If urus-server or axum-server is not
# requested, cargo will still build the workspace, but that's fine — a
# 1-minute one-time cost beats a confusing "binary not found" mid-batch.

echo "run_all: building binaries (release)..."
( cd "$REPO_ROOT" && cargo build --release ) || {
  echo "run_all: cargo build failed" >&2
  exit 3
}

# ---------------------------------------------------------------------------
# Run the matrix
# ---------------------------------------------------------------------------

total="${#runs[@]}"
echo
echo "run_all: batch=${BATCH}, ${total} run(s):"
i=0
for r in "${runs[@]}"; do
  i=$((i + 1))
  printf "  [%d/%d] %s\n" "$i" "$total" "$r"
done
echo

i=0
failed=()
succeeded=()
START_TS=$(date +%s)

# Per-run port offset so concurrent stale processes (if any) don't collide.
BASE_PORT="${BASE_PORT:-8080}"

for r in "${runs[@]}"; do
  i=$((i + 1))
  set -- $r
  SERVER="$1"; SCENARIO="$2"
  RUN_ID="${SERVER}-${SCENARIO}"
  RUN_DIR="${BATCH_DIR}/${RUN_ID}"

  printf '\n========== [%d/%d] %-12s %-3s ==========\n' \
    "$i" "$total" "$SERVER" "$SCENARIO"

  # Each run gets its own port so a hung server from a previous run can't
  # silently keep answering on the default port.
  PORT=$((BASE_PORT + i))

  # Each run gets a unique results dir under the batch dir. runner.sh
  # writes to results/<run_id>/ at the repo root by default — we pass the
  # run_id but also need it to land under the batch dir. Cheapest fix:
  # let runner write to its default location, then move it under batch.
  if PORT="$PORT" "${REPO_ROOT}/runner.sh" "$SERVER" "$SCENARIO" "$RUN_ID" \
       > "${BATCH_DIR}/${RUN_ID}.stdout.log" 2>&1; then
    # Move the runner's output dir into the batch dir.
    if [[ -d "${REPO_ROOT}/results/${RUN_ID}" ]]; then
      rm -rf "$RUN_DIR"
      mv "${REPO_ROOT}/results/${RUN_ID}" "$RUN_DIR"
    fi
    # Move the stdout log next to the run.
    mv "${BATCH_DIR}/${RUN_ID}.stdout.log" "${RUN_DIR}/run_all.log" 2>/dev/null || true

    # Pretty-print the headline line from the result.
    if [[ -f "${RUN_DIR}/result.json" ]]; then
      rps=$(awk -F'[:,]' '/"sustained_rps"/ {gsub(/[ "]/, "", $2); print $2; exit}' "${RUN_DIR}/result.json")
      p99=$(awk -F'"' '/"p99"/ {print $4; exit}' "${RUN_DIR}/result.json")
      printf '       sustained=%s rps   p99=%s\n' "$rps" "$p99"
    fi
    succeeded+=("$r")
  else
    echo "       FAILED (see ${BATCH_DIR}/${RUN_ID}.stdout.log)"
    failed+=("$r")
    # Also try to salvage anything runner left behind for debugging.
    if [[ -d "${REPO_ROOT}/results/${RUN_ID}" ]]; then
      mkdir -p "${BATCH_DIR}/${RUN_ID}"
      mv "${REPO_ROOT}/results/${RUN_ID}"/* "${BATCH_DIR}/${RUN_ID}/" 2>/dev/null || true
      rmdir "${REPO_ROOT}/results/${RUN_ID}" 2>/dev/null || true
    fi
  fi
done

END_TS=$(date +%s)
ELAPSED=$((END_TS - START_TS))

# ---------------------------------------------------------------------------
# Build summary.md
# ---------------------------------------------------------------------------
#
# The headline-table shape from urus-bench-spec.md §6 has servers as
# columns, scenarios as rows, two metrics per scenario (RPS sustained,
# p99). We render whatever subset of the matrix actually ran.

SUMMARY="${BATCH_DIR}/summary.md"
INDEX="${BATCH_DIR}/index.json"

# Discover every server / scenario actually represented in the batch.
declare -A by_run        # by_run["server scenario"] = path to result.json
servers_seen=()
scenarios_seen=()
for d in "$BATCH_DIR"/*/; do
  rj="${d}result.json"
  [[ -f "$rj" ]] || continue
  srv=$(awk -F'"' '/"server":/ {print $4; exit}' "$rj")
  scn=$(awk -F'"' '/"scenario":/ {print $4; exit}' "$rj")
  [[ -z "$srv" || -z "$scn" ]] && continue
  by_run["$srv $scn"]="$rj"
  # Track distinct lists (order preserved by first occurrence).
  if ! printf '%s\n' "${servers_seen[@]}" | grep -qx "$srv"; then
    servers_seen+=("$srv")
  fi
  if ! printf '%s\n' "${scenarios_seen[@]}" | grep -qx "$scn"; then
    scenarios_seen+=("$scn")
  fi
done

# Sort scenarios in canonical order s1..s4.
IFS=$'\n' scenarios_seen=($(printf '%s\n' "${scenarios_seen[@]}" | sort))
unset IFS

# ---------------------------------------------------------------------------
# Emit summary.md
# ---------------------------------------------------------------------------

{
  echo "# urus-bench summary"
  echo
  echo "- **Batch:** \`${BATCH}\`"
  echo "- **Host:** \`$(uname -n)\` (\`$(uname -srm)\`)"
  echo "- **CPUs:** \`$(nproc) cores\`"
  if [[ -r /proc/cpuinfo ]]; then
    model=$(awk -F: '/^model name/ {print $2; exit}' /proc/cpuinfo | sed 's/^ *//')
    [[ -n "$model" ]] && echo "- **Model:** ${model}"
  fi
  echo "- **Started:** \`$(date -d "@$START_TS" -Iseconds 2>/dev/null || date -Iseconds)\`"
  echo "- **Elapsed:** ${ELAPSED}s"
  echo "- **Successful runs:** ${#succeeded[@]} / ${total}"
  if [[ "${#failed[@]}" -gt 0 ]]; then
    echo "- **Failed runs:** ${failed[*]}"
  fi
  echo
  echo "Timings: probe=${PROBE_SEC:-30}s, warmup=${WARMUP_SEC:-30}s, measure=${MEASURE_SEC:-60}s, sat_ratio=${SAT_RATIO:-0.70}"
  echo

  if [[ "${#servers_seen[@]}" -eq 0 ]]; then
    echo "_No successful runs to summarize._"
    exit 0
  fi

  # --- Sustained RPS table -----------------------------------------------
  echo "## Sustained RPS (wrk2 @ ${SAT_RATIO:-0.70} of saturation)"
  echo
  printf '|            '
  for srv in "${servers_seen[@]}"; do printf '| %-14s ' "$srv"; done
  printf '|\n'
  printf '|------------'
  for _ in "${servers_seen[@]}"; do printf '|----------------'; done
  printf '|\n'
  for scn in "${scenarios_seen[@]}"; do
    printf '| %-10s ' "$scn"
    for srv in "${servers_seen[@]}"; do
      rj="${by_run["$srv $scn"]:-}"
      if [[ -n "$rj" ]]; then
        rps=$(awk -F'[:,]' '/"sustained_rps"/ {gsub(/[ "]/, "", $2); print $2; exit}' "$rj")
        # round to integer for readability
        rps=$(awk -v x="$rps" 'BEGIN {printf "%d", x+0.5}')
        printf '| %14s ' "$rps"
      else
        printf '| %14s ' "—"
      fi
    done
    printf '|\n'
  done
  echo

  # --- p99 table ----------------------------------------------------------
  echo "## p99 latency (at target rate)"
  echo
  printf '|            '
  for srv in "${servers_seen[@]}"; do printf '| %-14s ' "$srv"; done
  printf '|\n'
  printf '|------------'
  for _ in "${servers_seen[@]}"; do printf '|----------------'; done
  printf '|\n'
  for scn in "${scenarios_seen[@]}"; do
    printf '| %-10s ' "$scn"
    for srv in "${servers_seen[@]}"; do
      rj="${by_run["$srv $scn"]:-}"
      if [[ -n "$rj" ]]; then
        p99=$(awk -F'"' '/"p99"/ {print $4; exit}' "$rj")
        printf '| %14s ' "$p99"
      else
        printf '| %14s ' "—"
      fi
    done
    printf '|\n'
  done
  echo

  # --- Spec pass/fail check, urus vs axum --------------------------------
  # Per spec §6: urus RPS >= 50% of axum RPS on the same scenario.
  printf '## Spec checks (urus vs axum)\n\n'
  printf '| Scenario | urus RPS | axum RPS | urus/axum | within 2× of axum |\n'
  printf '|----------|----------|----------|-----------|-------------------|\n'
  for scn in "${scenarios_seen[@]}"; do
    # Match urus-{mem,sqlite} against axum-{mem,sqlite} for the matching tier.
    urus_run=""; axum_run=""
    for k in "${!by_run[@]}"; do
      set -- $k
      s="$1"; sc="$2"
      [[ "$sc" != "$scn" ]] && continue
      case "$s" in
        urus-*) urus_run="${by_run[$k]}" ;;
        axum-*) axum_run="${by_run[$k]}" ;;
      esac
    done
    if [[ -n "$urus_run" && -n "$axum_run" ]]; then
      u=$(awk -F'[:,]' '/"sustained_rps"/ {gsub(/[ "]/, "", $2); print $2; exit}' "$urus_run")
      a=$(awk -F'[:,]' '/"sustained_rps"/ {gsub(/[ "]/, "", $2); print $2; exit}' "$axum_run")
      ratio=$(awk -v u="$u" -v a="$a" 'BEGIN { if (a+0 > 0) printf "%.2f", (u+0)/(a+0); else print "n/a" }')
      pass=$(awk -v r="$ratio" 'BEGIN { print (r != "n/a" && r+0 >= 0.5) ? "✅" : "❌" }')
      printf '| %-8s | %8.0f | %8.0f | %9s | %-17s |\n' "$scn" "$u" "$a" "$ratio" "$pass"
    else
      printf '| %-8s | %8s | %8s | %9s | %-17s |\n' "$scn" "—" "—" "—" "—"
    fi
  done
  echo

  # --- per-run detail links ----------------------------------------------
  echo "## Per-run artefacts"
  echo
  for r in "${succeeded[@]}"; do
    set -- $r
    s="$1"; sc="$2"
    echo "- \`${s}-${sc}/\` — [result.json](${s}-${sc}/result.json), [measure.txt](${s}-${sc}/measure.txt), [server.log](${s}-${sc}/server.log)"
  done
  if [[ "${#failed[@]}" -gt 0 ]]; then
    echo
    echo "### Failed runs"
    for r in "${failed[@]}"; do
      set -- $r
      s="$1"; sc="$2"
      echo "- \`${s}-${sc}\` — see \`${s}-${sc}.stdout.log\`"
    done
  fi
} > "$SUMMARY"

# ---------------------------------------------------------------------------
# Emit machine-readable index.json
# ---------------------------------------------------------------------------

{
  echo "{"
  echo "  \"batch\": \"${BATCH}\","
  echo "  \"elapsed_seconds\": ${ELAPSED},"
  echo "  \"runs\": ["
  first=1
  for k in "${!by_run[@]}"; do
    set -- $k
    s="$1"; sc="$2"
    rj="${by_run[$k]}"
    [[ "$first" -eq 1 ]] || echo "    ,"
    first=0
    echo "    {"
    echo "      \"server\":   \"${s}\","
    echo "      \"scenario\": \"${sc}\","
    echo "      \"result_path\": \"${s}-${sc}/result.json\","
    # Inline the result.json contents, indented; simplest is to just paste.
    awk '/^\{/{p=1} p{print "      " $0} /^\}/{p=0}' "$rj" \
      | sed '1d' | head -n -1 | sed 's/^/      "result": {/; t; s/^/      /' >/dev/null
    # Simpler: just reference the file path. The result.json is already
    # right there. We don't need to inline.
    echo "    }"
  done
  echo "  ]"
  echo "}"
} > "$INDEX"

# ---------------------------------------------------------------------------
# Wrap up
# ---------------------------------------------------------------------------

echo
echo "=========================================================="
echo "run_all: DONE"
echo "  batch:      ${BATCH}"
echo "  elapsed:    ${ELAPSED}s"
echo "  successful: ${#succeeded[@]} / ${total}"
[[ "${#failed[@]}" -gt 0 ]] && echo "  failed:     ${failed[*]}"
echo "  summary:    ${SUMMARY}"
echo "=========================================================="
echo
cat "$SUMMARY"