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

# Reproducible larson runner for hakmem/system/mimalloc.
#
# Usage:
#   scripts/run_larson.sh [runtime_sec] [threads]
# Examples:
#   scripts/run_larson.sh                 # default: 10s, threads=1 4
#   scripts/run_larson.sh 10 1            # 10s, 1 thread
#
# Optional env vars:
#   HAKMEM_WRAP_TINY=0|1
#   HAKMEM_WRAP_TINY_REFILL=0|1
#   HAKMEM_TINY_MAG_CAP=INT
#   HAKMEM_SAFE_FREE=0|1
#   HAKMEM_EVO_SAMPLE=INT  (0 disables evo recording; default 0)
#   MIMALLOC_SO=/path/to/libmimalloc.so.2 (optional; if not set, auto-detect)

usage() {
  cat << USAGE
Usage: scripts/run_larson.sh [options] [runtime_sec] [threads_csv]

Options:
  -d SECONDS     Runtime seconds (default: 10)
  -t CSV         Threads CSV, e.g. 1,4 (default: 1,4)
  -c NUM         Chunks per thread (default: 10000)
  -r NUM         Rounds (default: 1)
  -m BYTES       Min size (default: 8)
  -M BYTES       Max size (default: 1024)
  -s SEED        Random seed (default: 12345)
  -p PRESET      Preset: burst|loop (sets -c/-r)
  -w             Include WRAP_TINY=1 runs (default: off)
  -h             Show this help

Env overrides (alternative to flags):
  MIN, MAX, CHUNK_PER_THREAD, ROUNDS, SEED
  HAKMEM_* toggles per README
USAGE
}

# Defaults
RUNTIME="10"
THREADS_ARG="1,4"

# Workload defaults (burst preset)
MIN="${MIN:-8}"
MAX="${MAX:-1024}"
CHUNK_PER_THREAD="${CHUNK_PER_THREAD:-10000}"
ROUNDS="${ROUNDS:-1}"
SEED="${SEED:-12345}"

PRESET=""

INCLUDE_WRAP=0
while getopts ":d:t:c:r:m:M:s:p:wh" opt; do
  case $opt in
    d) RUNTIME="$OPTARG" ;;
    t) THREADS_ARG="$OPTARG" ;;
    c) CHUNK_PER_THREAD="$OPTARG" ;;
    r) ROUNDS="$OPTARG" ;;
    m) MIN="$OPTARG" ;;
    M) MAX="$OPTARG" ;;
    s) SEED="$OPTARG" ;;
    p) PRESET="$OPTARG" ;;
    w) INCLUDE_WRAP=1 ;;
    h) usage; exit 0 ;;
    :) echo "Missing argument for -$OPTARG" >&2; usage; exit 2 ;;
    *) usage; exit 2 ;;
  esac
done
shift $((OPTIND-1))

# Backward-compatible positional args
if [[ $# -ge 1 ]]; then RUNTIME="$1"; fi
if [[ $# -ge 2 ]]; then THREADS_ARG="$2"; fi

case "$PRESET" in
  burst|BURST)
    CHUNK_PER_THREAD=10000; ROUNDS=1 ;;
  loop|LOOP)
    CHUNK_PER_THREAD=100; ROUNDS=100 ;;
  "" ) : ;;
  *) echo "Unknown preset: $PRESET" >&2; exit 2 ;;
esac

# Params matching our standard runs (larson reads: runtime, min, max, chunks/thread, rounds, seed, threads)
# Show resolved parameters for reproducibility
echo "[CFG] runtime=${RUNTIME}s threads={${THREADS_ARG}} min=${MIN} max=${MAX} chunks/thread=${CHUNK_PER_THREAD} rounds=${ROUNDS} seed=${SEED}"

ROOT_DIR="$(cd "$(dirname "$0")/.." && pwd)"
LIB_HAKMEM="$ROOT_DIR/libhakmem.so"
LARSON_BIN="$ROOT_DIR/mimalloc-bench/bench/larson/larson"

if [[ ! -x "$LARSON_BIN" ]]; then
  echo "[ERR] Larson binary not found at: $LARSON_BIN" >&2
  echo "      Did you sync submodule/build bench?" >&2
  exit 1
fi

if [[ ! -f "$LIB_HAKMEM" ]]; then
  echo "[INFO] libhakmem.so not found; building..."
  (cd "$ROOT_DIR" && make -j4 shared >/dev/null)
fi

abs_hakmem="$(readlink -f "$LIB_HAKMEM")"

detect_mimalloc() {
  if [[ -n "${MIMALLOC_SO:-}" && -f "$MIMALLOC_SO" ]]; then
    echo "$MIMALLOC_SO"
    return 0
  fi
  # try common paths or ldconfig
  for p in \
    /usr/lib/x86_64-linux-gnu/libmimalloc.so.2 \
    /lib/x86_64-linux-gnu/libmimalloc.so.2; do
    [[ -f "$p" ]] && { echo "$p"; return 0; }
  done
  if command -v ldconfig >/dev/null 2>&1; then
    so="$(ldconfig -p | awk '/libmimalloc.so/ {print $4; exit}')"
    [[ -n "$so" && -f "$so" ]] && { echo "$so"; return 0; }
  fi
  return 1
}

mimalloc_so=""
if mimalloc_so=$(detect_mimalloc); then
  :
else
  mimalloc_so=""
fi

run_case() {
  local label="$1"; shift
  local preload="$1"; shift
  local threads="$1"; shift

  echo "\n== $label | ${threads}T | ${RUNTIME}s =="
  if [[ -n "$preload" ]]; then
    env LD_PRELOAD="$preload" "$LARSON_BIN" "$RUNTIME" "$MIN" "$MAX" "$CHUNK_PER_THREAD" "$ROUNDS" "$SEED" "$threads" 2>&1 | tail -n 3
  else
    "$LARSON_BIN" "$RUNTIME" "$MIN" "$MAX" "$CHUNK_PER_THREAD" "$ROUNDS" "$SEED" "$threads" 2>&1 | tail -n 3
  fi
}

IFS=',' read -r -a THREADS <<< "$THREADS_ARG"

for t in "${THREADS[@]}"; do
  # system malloc
  run_case "system malloc" "" "$t"

  # mimalloc (optional)
  if [[ -n "$mimalloc_so" ]]; then
    run_case "mimalloc" "$mimalloc_so" "$t"
  else
    echo "\n== mimalloc | ${t}T | ${RUNTIME}s =="
    echo "[SKIP] libmimalloc not found"
  fi

  # hakmem default
  run_case "hakmem (default)" "$abs_hakmem" "$t"

  # hakmem wrap tiny (optional)
  if [[ "$INCLUDE_WRAP" -eq 1 ]]; then
    echo "\n== hakmem (HAKMEM_WRAP_TINY=1) | ${t}T | ${RUNTIME}s =="
    HAKMEM_WRAP_TINY=1 LD_PRELOAD="$abs_hakmem" "$LARSON_BIN" "$RUNTIME" "$MIN" "$MAX" "$CHUNK_PER_THREAD" "$ROUNDS" "$SEED" "$t" 2>&1 | tail -n 3
  fi
done

echo "\nDone."