#!/usr/bin/env bash
# Purpose: load one model in llama-server at several context lengths and KV cache types, one
#          after another, and record for each what llama.cpp allocated: the KV cache size it
#          logged, the per-device memory breakdown it prints on exit, and the KV size the
#          formula predicts from the model's own metadata
# Platform: all (Linux, macOS and WSL2; the backend is whatever llama.cpp was built with)
# Minimum memory: 8 GB
# Assumes: llama.cpp v0.4.0 or later built as in this part's install lesson (llama-server in
#          $LLAMA_BIN, on PATH, or in ~/llama.cpp/build/bin); a GGUF model file; curl, awk
#          and python3 on PATH; nothing listening on $PORT; nothing else large running
#
# Usage: bash kv-memory-ladder.sh <model.gguf> [labbook.md]
#
# Environment:
#   LLAMA_BIN     directory holding llama-server  (default: PATH, then ~/llama.cpp/build/bin)
#   LADDER        context:cache-type rungs, space-separated
#                 (default: "4096:f16 16384:f16 32768:f16 32768:q8_0")
#   NGL           layers to offload                (default: 999, meaning all)
#   PORT          local port for the probe         (default: 8089)
#   LOAD_TIMEOUT  seconds to wait for a load       (default: 600)
#   LOG_DIR       where each rung's full log goes  (default: ./kv-logs)
#
# Each rung starts the server bound to 127.0.0.1 with --verbose (the KV cache and memory
# breakdown lines are not printed at the default verbosity), waits for "listening on" or for
# the process to exit, then stops it with SIGTERM; the memory breakdown is printed on the way
# out. --fit off stops llama.cpp adjusting anything to fit, so a rung that does not fit fails
# instead of silently becoming a different rung. Nothing is sent to the server.

set -euo pipefail

MODEL="${1:-}"
LABBOOK="${2:-}"
LADDER="${LADDER:-4096:f16 16384:f16 32768:f16 32768:q8_0}"
NGL="${NGL:-999}"
PORT="${PORT:-8089}"
LOAD_TIMEOUT="${LOAD_TIMEOUT:-600}"
LOG_DIR="${LOG_DIR:-./kv-logs}"

die() { echo "kv-memory-ladder: $*" >&2; exit 1; }

[ -n "$MODEL" ] || die "usage: bash kv-memory-ladder.sh <model.gguf> [labbook.md]"
[ -f "$MODEL" ] || die "model file '$MODEL' does not exist (names are case-sensitive)"
for tool in curl awk python3; do
  command -v "$tool" >/dev/null || die "$tool is not installed or not on PATH"
done

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

if curl -s --max-time 2 "http://127.0.0.1:$PORT/health" >/dev/null 2>&1; then
  die "something is already answering on 127.0.0.1:$PORT; stop it or set PORT to a free port"
fi

mkdir -p "$LOG_DIR"
VERSION="$("$SERVER" --version 2>&1 | head -n 1)"
BASE="$(basename "$MODEL" .gguf)"
echo "==> $BASE with $VERSION"
printf '%-8s %-6s %-9s %12s %12s   %s\n' "context" "cache" "status" "KV MiB" "predicted" "device breakdown (total = free + (self = model + context + compute) + unaccounted)"

SERVER_PID=""
cleanup() { if [ -n "$SERVER_PID" ]; then kill "$SERVER_PID" 2>/dev/null || true; fi; }
trap cleanup EXIT
trap 'cleanup; exit 130' INT TERM

# The value of one "print_info: <key> = <value>" line of the current rung's log.
meta() { awk -v key="$1" '$0 ~ "print_info: " key " +=" { v = $NF } END { print v }' "$log"; }

for rung in $LADDER; do
  ctx="${rung%%:*}"
  cache="${rung##*:}"
  log="$LOG_DIR/${BASE}__c${ctx}__${cache}.log"

  "$SERVER" --model "$MODEL" --n-gpu-layers "$NGL" --ctx-size "$ctx" --parallel 1 \
    --flash-attn on --cache-type-k "$cache" --cache-type-v "$cache" --fit off \
    --host 127.0.0.1 --port "$PORT" --verbose > "$log" 2>&1 &
  SERVER_PID=$!

  status="timeout"
  waited=0
  while [ "$waited" -lt "$LOAD_TIMEOUT" ]; do
    if grep -q "listening on" "$log"; then status="loaded"; break; fi
    if ! kill -0 "$SERVER_PID" 2>/dev/null; then status="failed"; break; fi
    sleep 1
    waited=$((waited + 1))
  done
  kill -TERM "$SERVER_PID" 2>/dev/null || true
  wait "$SERVER_PID" 2>/dev/null || true
  SERVER_PID=""

  # "llama_kv_cache: size = 576.00 MiB (4096 cells, ...)"; a model with sliding-window
  # layers logs two caches, so the sizes are summed.
  kv_mib="$(awk '/llama_kv_cache: size =/ { for (i = 1; i <= NF; i++) if ($i == "=") { s += $(i + 1); break } } END { if (s > 0) printf "%.2f", s }' "$log")"
  predicted="$(python3 -c '
import sys
layers, heads, hk, hv, ctx, cache = sys.argv[1:7]
per_element = {"f32": 4.0, "f16": 2.0, "bf16": 2.0, "q8_0": 34 / 32, "q4_0": 18 / 32}.get(cache)
try:
    print(f"{int(layers) * int(heads) * (int(hk) + int(hv)) * per_element * int(ctx) / 2**20:.2f}")
except (TypeError, ValueError):
    print("")
' "$(meta n_layer)" "$(meta n_head_kv)" "$(meta n_embd_head_k)" "$(meta n_embd_head_v)" "$ctx" "$cache")"
  breakdown="$(grep 'common_memory_breakdown_print: |' "$log" | grep -v 'memory breakdown \[MiB\]' | sed 's/.*common_memory_breakdown_print: //' | tail -n 4 || true)"
  first_row="$(printf '%s\n' "$breakdown" | head -n 1)"
  printf '%-8s %-6s %-9s %12s %12s   %s\n' "$ctx" "$cache" "$status" "${kv_mib:--}" "${predicted:--}" "${first_row:--}"
  if [ "$status" != "loaded" ]; then
    grep -iE 'out of memory|OutOfDeviceMemory|failed to allocate|unable to allocate|error' "$log" | tail -n 3 | sed 's/^/           /' || true
  fi

  if [ -n "$LABBOOK" ]; then
    python3 -c '
import json, sys
from datetime import datetime, timezone
model, ctx, cache, status, kv, predicted, breakdown, version, log = sys.argv[1:10]
print(json.dumps({
    "lab": "part-06/kv-memory-ladder",
    "date": datetime.now(timezone.utc).isoformat(timespec="seconds"),
    "engine": "llama.cpp", "version": version, "model_path": model,
    "ctx_size": int(ctx), "cache_type": cache, "status": status,
    "kv_mib": float(kv) if kv else None,
    "kv_mib_predicted": float(predicted) if predicted else None,
    "memory_breakdown_mib": [line for line in breakdown.splitlines() if line.strip()],
    "log": log,
}))' "$MODEL" "$ctx" "$cache" "$status" "$kv_mib" "$predicted" "$breakdown" "$VERSION" "$log" >> "$LABBOOK"
  fi
done

echo "==> Full logs in $LOG_DIR${LABBOOK:+; one line per rung appended to $LABBOOK}"