#!/usr/bin/env bash
# Purpose: run llama-bench over a list of GGUF models with one fixed methodology, keep the
#          raw JSON and the error log for every model, and append one notebook line per test
# Platform: all (the same command on every track; only the backend differs)
# Minimum memory: 8 GB
# Assumes: llama.cpp built as in this part's install lesson (llama-bench found in $LLAMA_BIN,
#          on PATH, or in ~/llama.cpp/build/bin); models downloaded with Part 4's
#          fetch-model.sh; python3 for bench-to-labbook.py, which sits beside this script;
#          nothing else large running, because a benchmark shares the machine with it
#
# Usage: bash bench-reference-models.sh [labbook.md] [model.gguf ...]
#   With model files: benchmark exactly those, in the order given.
#   Without:          benchmark every *.gguf under $MODELS_DIR, sorted by path.
#   For a split model, give the first shard (...-00001-of-0000N.gguf); llama.cpp opens the
#   others from the same directory, and later shards are skipped if given.
#
# Environment:
#   MODELS_DIR   the model library, used when no files are given (default: $HOME/models)
#   LLAMA_BIN    directory holding llama-bench     (default: PATH, then ~/llama.cpp/build/bin)
#   RESULTS_DIR  where raw JSON and logs are kept  (default: ./bench-results)
#   PROMPT_LEN   prompt-processing test, tokens     (default: 512)
#   GEN_LEN      generation test, tokens            (default: 128)
#   REPS         repetitions per test               (default: 5)
#   NGL          layers to offload                  (default: 999, more than any model has)
#   FLASH_ATTN   on, off or auto                    (default: on)
#   GGML_CUDA_ENABLE_UNIFIED_MEMORY  not read here; when set, llama.cpp's CUDA and HIP
#                backends read it, and bench-to-labbook.py records its value on every line
#
# One llama-bench run per model, two tests per run: prompt processing at PROMPT_LEN tokens
# and generation at GEN_LEN tokens. llama-bench runs one untimed warm-up of each test, then
# REPS timed repetitions, and reports the mean and standard deviation of the per-repetition
# rates. The options are the long forms listed by `llama-bench --help` at llama.cpp v0.4.0.
# A result that bench-to-labbook.py refuses (the CPU backend) is counted as not recorded; its
# raw JSON is kept.

set -euo pipefail

LABBOOK="${1:-labbook.md}"
if [ "$#" -gt 0 ]; then shift; fi
MODELS_DIR="${MODELS_DIR:-$HOME/models}"
RESULTS_DIR="${RESULTS_DIR:-./bench-results}"
PROMPT_LEN="${PROMPT_LEN:-512}"
GEN_LEN="${GEN_LEN:-128}"
REPS="${REPS:-5}"
NGL="${NGL:-999}"
FLASH_ATTN="${FLASH_ATTN:-on}"

die() { echo "bench-reference-models: $*" >&2; exit 1; }

case "$LABBOOK" in
  *.gguf) die "the first argument is the notebook; usage: bash bench-reference-models.sh labbook.md model.gguf ..." ;;
esac

if [ -n "${LLAMA_BIN:-}" ]; then
  BENCH="$LLAMA_BIN/llama-bench"
elif command -v llama-bench >/dev/null 2>&1; then
  BENCH="$(command -v llama-bench)"
else
  BENCH="$HOME/llama.cpp/build/bin/llama-bench"
fi
[ -x "$BENCH" ] || die "llama-bench not found at $BENCH; set LLAMA_BIN to the directory holding it"
command -v python3 >/dev/null || die "python3 is not installed"

SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
TO_LABBOOK="$SCRIPT_DIR/bench-to-labbook.py"
[ -f "$TO_LABBOOK" ] || die "bench-to-labbook.py is not beside this script in $SCRIPT_DIR"

# --- 1. Collect the models -------------------------------------------------------------
CANDIDATES=()
if [ "$#" -gt 0 ]; then
  for f in "$@"; do
    [ -f "$f" ] || die "model file '$f' does not exist (names are case-sensitive)"
    CANDIDATES+=("$f")
  done
else
  [ -d "$MODELS_DIR" ] || die "$MODELS_DIR does not exist; download models with fetch-model.sh first"
  while IFS= read -r f; do
    CANDIDATES+=("$f")
  done < <(find "$MODELS_DIR" -type f -name '*.gguf' | sort)
fi

[ "${#CANDIDATES[@]}" -gt 0 ] || die "no .gguf files found under $MODELS_DIR"

# Empty arrays are guarded before every "${array[@]}": macOS ships bash 3.2, where expanding
# an empty array under set -u is an error.
MODELS=()
for f in "${CANDIDATES[@]}"; do
  base="$(basename "$f")"
  case "$base" in
    mmproj*|*mmproj*)
      echo "    skipping $base (a multimodal projector, not a language model)"; continue ;;
    *-0000[2-9]-of-*|*-000[1-9][0-9]-of-*)
      echo "    skipping $base (a later shard; the first shard loads it)"; continue ;;
  esac
  MODELS+=("$f")
done
[ "${#MODELS[@]}" -gt 0 ] || die "no model files to benchmark"

mkdir -p "$RESULTS_DIR"
RUN_ID="$(date -u +%Y%m%dT%H%M%SZ)"
HOSTINFO="$(uname -s)-$(uname -m)"

echo "==> ${#MODELS[@]} model file(s) to benchmark with $BENCH"
echo "    prompt processing ${PROMPT_LEN} tokens, generation ${GEN_LEN} tokens, ${REPS} repetitions"
echo "    --n-gpu-layers $NGL, --flash-attn $FLASH_ATTN, run $RUN_ID"

# --- 2. Benchmark each model -----------------------------------------------------------
OK=0
FAILED_NAMES=()
for model in "${MODELS[@]}"; do
  base="$(basename "$model" .gguf)"
  out="$RESULTS_DIR/${RUN_ID}__${base}.json"
  echo "==> $base"
  if "$BENCH" --model "$model" --n-prompt "$PROMPT_LEN" --n-gen "$GEN_LEN" \
       --repetitions "$REPS" --n-gpu-layers "$NGL" --flash-attn "$FLASH_ATTN" \
       --output json > "$out" 2> "$out.log"; then
    if python3 "$TO_LABBOOK" --input "$out" --labbook "$LABBOOK" \
         --run-id "$RUN_ID" --host "$HOSTINFO" --model-path "$model"; then
      OK=$((OK + 1))
    else
      FAILED_NAMES+=("$base")
      echo "    NOT RECORDED; the raw result stays in $out" >&2
    fi
  else
    FAILED_NAMES+=("$base")
    echo "    FAILED; the last lines of $out.log were:" >&2
    tail -n 5 "$out.log" >&2 || true
    if grep -qiE 'out of memory|OutOfDeviceMemory|failed to allocate|alloc.*failed|unable to allocate' "$out.log"; then
      echo "    That is an allocation failure: the model does not fit. Record it and carry on." >&2
    else
      echo "    No allocation error in the log: check the file with Part 4's verify-library.sh" >&2
      echo "    (for a split model, check its shards from the repository directory, as the page shows)." >&2
    fi
  fi
done

echo "==> Done: $OK recorded, ${#FAILED_NAMES[@]} failed or not recorded. Raw JSON and logs in $RESULTS_DIR."
if [ "${#FAILED_NAMES[@]}" -gt 0 ]; then
  printf '    did not complete: %s\n' "${FAILED_NAMES[@]}" >&2
fi
[ "$OK" -gt 0 ] || exit 1