#!/usr/bin/env bash
set -euo pipefail

# Run a small suite across system/mimalloc/hakmem and save logs under docs/benchmarks/<timestamp>.
# Focus: mimalloc-bench/bench/larson patterns that cover Tiny/Mid/Large/Big.
# Optionally include a 'hakmem best' run (WRAP L1 + learning on + DYN auto).

RUNTIME=${RUNTIME:-1}
THREADS_CSV=${THREADS:-"1,4"}
BEST=${BEST:-0}
# Hard wall-clock timeout per run (external). Defaults to RUNTIME+3s.
BENCH_TIMEOUT=${BENCH_TIMEOUT:-}
KILL_GRACE=${KILL_GRACE:-2}
if [[ -z "${BENCH_TIMEOUT}" ]]; then
  # Add small cushion to allow larson to exit cleanly
  BENCH_TIMEOUT=$(( RUNTIME + 3 ))
fi

ROOT_DIR="$(cd "$(dirname "$0")/.." && pwd)"
OUTDIR="$ROOT_DIR/docs/benchmarks/$(date +%Y%m%d_%H%M%S)_SUITE"
mkdir -p "$OUTDIR"

LARSON="$ROOT_DIR/mimalloc-bench/bench/larson/larson"
LIB_HAK="$ROOT_DIR/libhakmem.so"
LIB_HAK_ABS=$(readlink -f "$LIB_HAK")
LIB_MI="/lib/x86_64-linux-gnu/libmimalloc.so.2"

echo "Suite: RUNTIME=${RUNTIME}s THREADS=${THREADS_CSV} BEST=${BEST}" | tee "$OUTDIR/summary.txt"

run_case() {
  local label="$1"; shift
  local min="$1"; local max="$2"; shift 2
  local threads_csv="$1"; shift
  IFS=',' read -r -a TARR <<< "$threads_csv"
  for t in "${TARR[@]}"; do
    echo "== ${label} | ${t}T | system ==" | tee -a "$OUTDIR/summary.txt"
    timeout -k "${KILL_GRACE}s" "${BENCH_TIMEOUT}s" \
      "$LARSON" "$RUNTIME" "$min" "$max" 10000 1 12345 "$t" 2>&1 \
      | tee "$OUTDIR/${label}_system_T${t}.log" | tail -n 3 | tee -a "$OUTDIR/summary.txt"
    if [[ -f "$LIB_MI" ]]; then
      echo "== ${label} | ${t}T | mimalloc ==" | tee -a "$OUTDIR/summary.txt"
      timeout -k "${KILL_GRACE}s" "${BENCH_TIMEOUT}s" \
        env LD_PRELOAD="$LIB_MI" "$LARSON" "$RUNTIME" "$min" "$max" 10000 1 12345 "$t" 2>&1 \
        | tee "$OUTDIR/${label}_mimalloc_T${t}.log" | tail -n 3 | tee -a "$OUTDIR/summary.txt"
    fi
    if [[ -f "$LIB_HAK" ]]; then
      echo "== ${label} | ${t}T | hakmem(default) ==" | tee -a "$OUTDIR/summary.txt"
      timeout -k "${KILL_GRACE}s" "${BENCH_TIMEOUT}s" \
        env LD_PRELOAD="$LIB_HAK_ABS" "$LARSON" "$RUNTIME" "$min" "$max" 10000 1 12345 "$t" 2>&1 \
        | tee "$OUTDIR/${label}_hakmem_default_T${t}.log" | tail -n 3 | tee -a "$OUTDIR/summary.txt"
      if [[ "$BEST" == "1" ]]; then
        echo "== ${label} | ${t}T | hakmem(best) ==" | tee -a "$OUTDIR/summary.txt"
        timeout -k "${KILL_GRACE}s" "${BENCH_TIMEOUT}s" \
          env HAKMEM_WRAP_L2=1 HAKMEM_WRAP_L25=1 HAKMEM_LEARN=1 HAKMEM_DYN1_AUTO=1 HAKMEM_DYN2_AUTO=1 HAKMEM_HIST_SAMPLE=7 HAKMEM_WMAX_LEARN=1 HAKMEM_WMAX_DWELL_SEC=2 \
          LD_PRELOAD="$LIB_HAK_ABS" "$LARSON" "$RUNTIME" "$min" "$max" 10000 1 12345 "$t" 2>&1 \
          | tee "$OUTDIR/${label}_hakmem_best_T${t}.log" | tail -n 3 | tee -a "$OUTDIR/summary.txt"
      fi
    fi
  done
}

# Tiny band (8–64B)
run_case tiny 8 64 "$THREADS_CSV"
# Mid band (2–32KiB)
run_case mid 2048 32768 "$THREADS_CSV"
# Large band (64KiB–1MiB)
run_case large 65536 1048576 "$THREADS_CSV"
# Big band (2–4MiB)
run_case big 2097152 4194304 "$THREADS_CSV"

echo "Saved suite: $OUTDIR" | tee -a "$OUTDIR/summary.txt"