#!/usr/bin/env bash
# Reproducible ETL throughput benchmark.
#
# Three configurations against a fixed BENCH_DOCS slice of the Wikipedia
# 2003-05-16 cur dump:
#
#   serial            single process, single SQLite (Phase 1a)
#   parallel-shared   N shards, one shared SQLite (Phase 1b — WAL serialized)
#   attached          N shards, per-shard SQLite (Phase 2 — true parallel writes)
#
# Output: per-config wall, throughput (docs/s), and a one-line summary table.
set -uo pipefail

BENCH_DOCS=${1:-5000}
BENCH_DIR=${BENCH_DIR:-/tmp/arborist-bench}
SHARDS=${SHARDS:-4}
DUMP=${DUMP:-data/20030516_cur_tablesql.bz2}
ARBORIST=${ARBORIST:-.venv/bin/arborist}

mkdir -p "$BENCH_DIR"

if [[ ! -f "$DUMP" ]]; then
  echo "missing dump $DUMP — run 'make fetch-cur' first" >&2
  exit 2
fi

run_silent() {
  # $1 = label (writes wall time in seconds to stdout)
  local start end
  start=$(date +%s.%N)
  "$@" >/dev/null 2>&1
  end=$(date +%s.%N)
  echo "$end - $start" | bc -l
}

# Each row will be: label,wall_s,docs/s
declare -a results

bench_serial() {
  local db=$BENCH_DIR/serial.db
  rm -f "$db" "$db-"*
  local wall
  wall=$(run_silent "$ARBORIST" --db "$db" ingest --quiet \
           --source wikipedia_cur --path "$DUMP" --limit "$BENCH_DOCS")
  local docs
  docs=$(sqlite3 "$db" "SELECT COUNT(*) FROM documents")
  printf "serial,%s,%s\n" "$wall" "$docs"
}

bench_parallel_shared() {
  local db=$BENCH_DIR/shared.db
  rm -f "$db" "$db-"*
  local per=$((BENCH_DOCS / SHARDS))
  local start end wall
  start=$(date +%s.%N)
  for i in $(seq 0 $((SHARDS - 1))); do
    "$ARBORIST" --db "$db" ingest --quiet \
      --source wikipedia_cur --path "$DUMP" \
      --shard "$i/$SHARDS" --limit "$per" >/dev/null 2>&1 &
  done
  wait
  end=$(date +%s.%N)
  wall=$(echo "$end - $start" | bc -l)
  local docs
  docs=$(sqlite3 "$db" "SELECT COUNT(*) FROM documents")
  printf "parallel-shared,%s,%s\n" "$wall" "$docs"
}

bench_attached() {
  local dir=$BENCH_DIR/attached
  rm -rf "$dir"
  mkdir -p "$dir"
  local per=$((BENCH_DOCS / SHARDS))
  local start end wall
  start=$(date +%s.%N)
  for i in $(seq 0 $((SHARDS - 1))); do
    "$ARBORIST" ingest --quiet \
      --source wikipedia_cur --path "$DUMP" \
      --shards-dir "$dir" --shard "$i/$SHARDS" --limit "$per" >/dev/null 2>&1 &
  done
  wait
  end=$(date +%s.%N)
  wall=$(echo "$end - $start" | bc -l)
  local docs
  docs=$("$ARBORIST" --shards-dir "$dir" stats 2>/dev/null \
           | python3 -c "import json,sys; print(json.load(sys.stdin)['documents_total'])")
  printf "attached,%s,%s\n" "$wall" "$docs"
}

echo "=== arborist ETL benchmark — $BENCH_DOCS docs target, $SHARDS shards ==="
echo

results+=("$(bench_serial)")
results+=("$(bench_parallel_shared)")
results+=("$(bench_attached)")

# Pretty table.
printf "  %-18s %10s %10s %12s\n" "config" "wall_s" "docs" "docs/s"
printf "  %-18s %10s %10s %12s\n" "------" "------" "----" "------"
for row in "${results[@]}"; do
  IFS=',' read -r label wall docs <<<"$row"
  rate=$(echo "scale=1; $docs / $wall" | bc -l)
  printf "  %-18s %10.2f %10s %12s\n" "$label" "$wall" "$docs" "$rate"
done
echo
echo "  (raw CSV: $BENCH_DIR/results.csv)"
{ printf "config,wall_s,docs\n"; printf "%s\n" "${results[@]}"; } > "$BENCH_DIR/results.csv"