#!/usr/bin/env bash
#
# Benchmark a running dragoman across concurrency levels, so performance can be
# compared release to release.
#
# Samples DOIs from the store's pid_records (random rows, Crossref and
# DataCite works only), turns them into URLs of one kind, and runs oha at each
# concurrency level for a fixed duration. Each level's full oha report is kept
# as JSON, and one line per level is appended to a TSV that accumulates runs:
#
#   label  kind  concurrency  requests/s  p50 ms  p99 ms  success  status codes
#
# Run it on the server itself: requests from this machine that name no client
# are not rate limited (see "Request limits" in the README), so the numbers
# measure dragoman rather than its limits. Needs sqlite3, python3 and oha
# (`cargo install oha`).
#
# Usage:
#   scripts/bench.sh [options]
#
# Options (defaults in brackets):
#   -d PATH     commonmeta store to sample [$COMMONMETA_DB, else /var/lib/commonmeta/commonmeta.sqlite3]
#   -u URL      base URL of the server [https://127.0.0.1]
#   -k KIND     redirect | handles | page | bibtex [redirect]
#                 redirect  /{doi}                 302 to the landing page
#                 handles   /api/handles/{doi}     the handle proxy's answer
#                 page      /{doi}?noredirect      the entity page
#                 bibtex    /{doi}?format=bibtex   a format conversion
#   -n N        DOIs to sample [10000]
#   -c LEVELS   concurrency levels, space-separated ["1 2 4 8 16 32 64"]
#   -z TIME     duration of each level [30s]
#   -l LABEL    label for this run [the version `dragoman --version` reports]
#   -o DIR      where reports go [bench/]
#
# Example — compare two releases:
#   scripts/bench.sh -l 0.5.3          # before upgrading
#   scripts/bench.sh -l 0.5.4          # after
#   column -t bench/results.tsv

set -euo pipefail

DB="${COMMONMETA_DB:-/var/lib/commonmeta/commonmeta.sqlite3}"
BASE="https://127.0.0.1"
KIND="redirect"
SAMPLE=10000
LEVELS="1 2 4 8 16 32 64"
DURATION="30s"
LABEL=""
OUT="bench"

while getopts "d:u:k:n:c:z:l:o:h" opt; do
  case "$opt" in
    d) DB="$OPTARG" ;;
    u) BASE="${OPTARG%/}" ;;
    k) KIND="$OPTARG" ;;
    n) SAMPLE="$OPTARG" ;;
    c) LEVELS="$OPTARG" ;;
    z) DURATION="$OPTARG" ;;
    l) LABEL="$OPTARG" ;;
    o) OUT="$OPTARG" ;;
    h) sed -n '2,/^$/p' "$0" | sed 's/^# \{0,1\}//'; exit 0 ;;
    *) exit 2 ;;
  esac
done

for tool in sqlite3 python3 oha; do
  command -v "$tool" >/dev/null || { echo "error: $tool not found" >&2; exit 1; }
done
[[ -r "$DB" ]] || { echo "error: cannot read store '$DB'" >&2; exit 1; }
case "$KIND" in
  redirect) template='%s/%s' ;;
  handles) template='%s/api/handles/%s' ;;
  page) template='%s/%s?noredirect' ;;
  bibtex) template='%s/%s?format=bibtex' ;;
  *) echo "error: unknown kind '$KIND'" >&2; exit 2 ;;
esac
if [[ -z "$LABEL" ]]; then
  LABEL="$(dragoman --version 2>/dev/null | awk '{print $2}')"
  LABEL="${LABEL:-unknown}"
fi

mkdir -p "$OUT"
stamp="$(date -u +%Y%m%dT%H%M%SZ)"
run="$OUT/$stamp-$LABEL-$KIND"
mkdir -p "$run"

# Random rows by rowid: ORDER BY random() would scan the whole table. Draw
# three times the sample, since some rowids are gaps or not works.
echo "sampling $SAMPLE DOIs from $DB" >&2
sqlite3 -readonly "$DB" "
  WITH RECURSIVE
    top(n) AS (SELECT max(id) FROM pid_records),
    draw(i, id) AS (
      SELECT 1, abs(random()) % (SELECT n FROM top) + 1
      UNION ALL
      SELECT i + 1, abs(random()) % (SELECT n FROM top) + 1 FROM draw WHERE i < $((SAMPLE * 3))
    )
  SELECT DISTINCT p.pid FROM draw JOIN pid_records p ON p.id = draw.id
  WHERE p.source_id IN (1, 2)
  LIMIT $SAMPLE;" >"$run/dois.txt"
got="$(wc -l <"$run/dois.txt" | tr -d ' ')"
[[ "$got" -gt 0 ]] || { echo "error: sampled no DOIs" >&2; exit 1; }
echo "sampled $got DOIs" >&2
while read -r doi; do
  # shellcheck disable=SC2059
  printf "$template\n" "$BASE" "$doi"
done <"$run/dois.txt" >"$run/urls.txt"

results="$OUT/results.tsv"
[[ -s "$results" ]] || printf 'label\tkind\tconcurrency\trps\tp50_ms\tp99_ms\tsuccess\tstatus\tstarted\n' >"$results"

for c in $LEVELS; do
  echo "c=$c for $DURATION" >&2
  # --redirect 0: a redirect is the answer being measured, not a hop to follow.
  oha --no-tui --insecure --redirect 0 --urls-from-file "$run/urls.txt" \
    -c "$c" -z "$DURATION" --output-format json >"$run/c$c.json"
  python3 - "$run/c$c.json" "$LABEL" "$KIND" "$c" "$stamp" >>"$results" <<'EOF'
import json, sys
path, label, kind, c, stamp = sys.argv[1:]
r = json.load(open(path))
summary = r["summary"]
pct = r.get("latencyPercentiles", {})
ms = lambda v: f"{v * 1000:.1f}" if v is not None else ""
codes = ",".join(f"{k}:{v}" for k, v in sorted(r.get("statusCodeDistribution", {}).items()))
print("\t".join([
    label, kind, c,
    f"{summary['requestsPerSec']:.0f}",
    ms(pct.get("p50")), ms(pct.get("p99")),
    f"{summary['successRate'] * 100:.2f}%",
    codes, stamp,
]))
EOF
  tail -1 "$results" | cut -f3-7 >&2
done

echo "reports in $run, summary in $results" >&2