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

# Run (nearly) the full mimalloc-bench suite with timeouts.
# Compares: HAKMEM (via LD_PRELOAD on sys), mimalloc (mi), and system (sys).
#
# Env knobs:
#   TIMEOUT_SEC   per-run timeout seconds (default: 900)
#   PROCS         concurrency list for bench.sh --procs (default: 1,4)
#   INCLUDE_JE    include jemalloc reference (0/1, default: 0)

ROOT_DIR=$(cd "$(dirname "$0")/.." && pwd)
cd "$ROOT_DIR"

TIMEOUT_SEC=${TIMEOUT_SEC:-900}
PROCS=${PROCS:-1,4}
INCLUDE_JE=${INCLUDE_JE:-0}
TESTS=${TESTS:-allt}           # Space-separated tests or 'allt'
REPEATS=${REPEATS:-1}          # Pass through to bench.sh (-r)

RESULT_DIR="$ROOT_DIR/bench_results/mimalloc_full_$(date +%Y%m%d_%H%M%S)"
mkdir -p "$RESULT_DIR"

BENCH_ROOT="$ROOT_DIR/mimalloc-bench"
BENCH_OUT="$BENCH_ROOT/out/bench"

if [[ ! -d "$BENCH_OUT" ]]; then
  echo "[warn] mimalloc-bench/out/bench not found. Attempting auto-build (bench only)."
  if [[ -x "$BENCH_ROOT/build-bench-env.sh" ]]; then
    pushd "$BENCH_ROOT" >/dev/null
    ./build-bench-env.sh bench
    popd >/dev/null
  else
    echo "[error] build-bench-env.sh not found under mimalloc-bench."
    echo "        Please build manually: cd mimalloc-bench && ./build-bench-env.sh bench"
    exit 1
  fi
fi

if [[ ! -d "$BENCH_OUT" ]]; then
  echo "[error] mimalloc-bench/out/bench still missing after auto-build."
  echo "        Try: cd mimalloc-bench && ./build-bench-env.sh all"
  exit 1
fi

echo "[info] Building HAKMEM shared library with PGO for LD_PRELOAD"
make -s pgo-profile-shared pgo-build-shared >/dev/null

pushd "$BENCH_OUT" >/dev/null

run_case() {
  local name="$1"; shift
  local preload="$1"; shift
  local args=("$@")
  local log="$RESULT_DIR/${name}.log"
  echo "[case] $name | timeout=${TIMEOUT_SEC}s | args: ${args[*]}" | tee -a "$log"
  if [[ -n "$preload" ]]; then
    LD_PRELOAD="$preload" timeout -s INT "$TIMEOUT_SEC" bash ../../bench.sh -r="$REPEATS" "${args[@]}" 2>&1 | tee -a "$log" || true
  else
    timeout -s INT "$TIMEOUT_SEC" bash ../../bench.sh -r="$REPEATS" "${args[@]}" 2>&1 | tee -a "$log" || true
  fi
  # Save benchres.csv for this case if present
  if [[ -f benchres.csv ]]; then
    cp -f benchres.csv "$RESULT_DIR/${name}_benchres.csv" || true
  fi
}

if [[ "$TESTS" == "allt" ]]; then
  # HAKMEM vs mimalloc vs system (and optional jemalloc) for full-all tests
  run_case "hakmem_procs=${PROCS//,/}" "$ROOT_DIR/libhakmem.so" --procs="$PROCS" sys allt
  run_case "mimalloc_procs=${PROCS//,/}" "" --procs="$PROCS" mi allt
  run_case "system_procs=${PROCS//,/}" "" --procs="$PROCS" sys allt
  if [[ "$INCLUDE_JE" == "1" ]]; then
    run_case "jemalloc_procs=${PROCS//,/}" "" --procs="$PROCS" je allt
  fi
else
  # Split per test to enforce per-test timeouts and partial progress
  for t in $TESTS; do
    run_case "hakmem_${t}_p=${PROCS//,/}" "$ROOT_DIR/libhakmem.so" --procs="$PROCS" sys "$t"
    run_case "mimalloc_${t}_p=${PROCS//,/}" "" --procs="$PROCS" mi "$t"
    run_case "system_${t}_p=${PROCS//,/}" "" --procs="$PROCS" sys "$t"
    if [[ "$INCLUDE_JE" == "1" ]]; then
      run_case "jemalloc_${t}_p=${PROCS//,/}" "" --procs="$PROCS" je "$t"
    fi
  done
fi

popd >/dev/null

echo "[info] Logs: $RESULT_DIR"
echo "[hint] Parse logs to CSV: scripts/parse_mimalloc_logs.py $RESULT_DIR > $RESULT_DIR/summary.csv"