#!/bin/sh
# GPU Fit benchmark: measures how fast a model runs in your local Ollama and, if you say yes,
# shares the result at https://gpufit.dev so others can see real speeds for their hardware.
#
#   curl -fsSL https://gpufit.dev/bench.sh | sh                 # pick from your installed models
#   curl -fsSL https://gpufit.dev/bench.sh | sh -s -- qwen3:8b  # benchmark one model
#
# Options: --yes (send without asking), --no-send (only print), --host URL (default $OLLAMA_HOST
# or http://127.0.0.1:11434). Needs curl and Ollama. Sends: GPU name, GPU memory, GPU count,
# OS name, Ollama version, model tag, quantization, how much of the model was on the GPU,
# context length and measured speeds. Nothing else: no hostname, username or files.
set -eu

SITE="${GPUFIT_SITE:-https://gpufit.dev}"
HOST="${OLLAMA_HOST:-http://127.0.0.1:11434}"
MODEL="" SEND="ask" RUNS=3 TOKENS=256 CTX=4096
while [ $# -gt 0 ]; do
  case "$1" in
    --yes|-y) SEND="yes" ;;
    --no-send|-n) SEND="no" ;;
    --host) HOST="$2"; shift ;;
    -h|--help) sed -n '2,12p' "$0" 2>/dev/null || true; exit 0 ;;
    *) MODEL="$1" ;;
  esac
  shift
done
case "$HOST" in http*) ;; *) HOST="http://$HOST" ;; esac
HOST="${HOST%/}"

say() { printf '%s\n' "$*" >&2; }
die() { say "Error: $*"; exit 1; }
# Pull one field out of a flat JSON response (no jq needed).
jnum() { printf '%s' "$1" | tr -d '\n' | grep -o "\"$2\": *[0-9.eE+-]*" | head -n1 | sed 's/^[^:]*: *//'; }
jstr() { printf '%s' "$1" | tr -d '\n' | grep -o "\"$2\": *\"[^\"]*\"" | head -n1 | sed 's/^[^:]*: *"//; s/"$//'; }
clean() { printf '%s' "$1" | tr -d '"\\\000-\037' | cut -c1-80; }
TTY=""; if [ -r /dev/tty ] && (exec </dev/tty) 2>/dev/null; then TTY=1; fi
ask() { if [ -n "$TTY" ]; then printf '%s' "$1" >&2; read -r REPLY </dev/tty || REPLY=""; else REPLY=""; fi; }

command -v curl >/dev/null 2>&1 || die "curl is required."
VERSION_JSON=$(curl -fsS "$HOST/api/version" 2>/dev/null) || die "Ollama isn't answering at $HOST. Start it (open the Ollama app, or run 'ollama serve') and try again."
OLLAMA_VERSION=$(jstr "$VERSION_JSON" version)

# ---- choose a model ----
if [ -z "$MODEL" ]; then
  TAGS=$(curl -fsS "$HOST/api/tags") || die "Couldn't list your Ollama models."
  LIST=$(printf '%s' "$TAGS" | tr -d '\n' | grep -o '"name": *"[^"]*"' | sed 's/^"name": *"//; s/"$//' | grep -v -i 'embed' || true)
  [ -n "$LIST" ] || die "You have no models in Ollama yet. Try: ollama pull qwen3:8b"
  say "Your Ollama models:"
  printf '%s\n' "$LIST" | awk '{ printf "  %d) %s\n", NR, $0 }' >&2
  ask "Benchmark which one? [1] "
  PICK="${REPLY:-1}"
  case "$PICK" in *[!0-9]*) MODEL="$PICK" ;; *) MODEL=$(printf '%s\n' "$LIST" | sed -n "${PICK}p") ;; esac
  [ -n "$MODEL" ] || die "No model picked."
fi

SHOW=$(curl -fsS "$HOST/api/show" -d "{\"model\":\"$MODEL\"}" 2>/dev/null) || die "Ollama doesn't have '$MODEL'. Pull it first: ollama pull $MODEL"
QUANT=$(jstr "$SHOW" quantization_level)
[ -n "$QUANT" ] || QUANT="unknown"

# ---- detect the GPU ----
OS=$(uname -s | tr 'A-Z' 'a-z'); GPU="" MEM="" COUNT=1
if [ "$OS" = "darwin" ]; then
  GPU=$(sysctl -n machdep.cpu.brand_string 2>/dev/null || true)
  B=$(sysctl -n hw.memsize 2>/dev/null || echo 0); MEM=$(awk -v b="$B" 'BEGIN { printf "%.0f", b / 1073741824 }')
  CORES=$(system_profiler SPDisplaysDataType 2>/dev/null | awk -F': ' '/Total Number of Cores/ { print $2; exit }' || true)
  OS="macos"
elif command -v nvidia-smi >/dev/null 2>&1 && nvidia-smi -L >/dev/null 2>&1; then
  LINES=$(nvidia-smi --query-gpu=name,memory.total --format=csv,noheader,nounits)
  COUNT=$(printf '%s\n' "$LINES" | grep -c . || true)
  GPU=$(printf '%s\n' "$LINES" | head -n1 | cut -d, -f1)
  MEM=$(printf '%s\n' "$LINES" | head -n1 | cut -d, -f2 | awk '{ printf "%.1f", $1 / 1024 }')
elif command -v rocm-smi >/dev/null 2>&1; then
  GPU=$(rocm-smi --showproductname 2>/dev/null | grep -i 'card series' | head -n1 | sed 's/.*: *//')
  COUNT=$(rocm-smi --showproductname 2>/dev/null | grep -ci 'card series' || echo 1)
  B=$(rocm-smi --showmeminfo vram 2>/dev/null | grep -i 'total memory' | head -n1 | grep -o '[0-9]*$' || echo 0)
  MEM=$(awk -v b="${B:-0}" 'BEGIN { if (b > 0) printf "%.1f", b / 1073741824 }')
elif command -v lspci >/dev/null 2>&1; then
  GPU=$(lspci | grep -i -E 'vga|3d controller' | grep -v -i 'intel' | head -n1 | sed 's/.*: //')
fi
GPU=$(clean "$GPU"); [ -n "$GPU" ] || GPU="unknown"
[ -z "${CORES:-}" ] || say "Detected $GPU with a $CORES-core GPU."

# ---- run ----
PROMPT="Write a detailed, practical guide to growing tomatoes on a balcony. Cover choosing a variety, containers, soil, watering, feeding, pruning, pests and harvesting."
gen() {
  curl -fsS "$HOST/api/generate" -d "{\"model\":\"$MODEL\",\"prompt\":\"$PROMPT\",\"stream\":false,\"options\":{\"num_predict\":$1,\"num_ctx\":$CTX,\"temperature\":0,\"seed\":42}}"
}
say "Loading $MODEL ($QUANT)..."
gen 1 >/dev/null || die "Ollama couldn't run $MODEL."
PS=$(curl -fsS "$HOST/api/ps" 2>/dev/null | tr -d '\n' || true)
# With several models loaded, read the sizes from our model's entry.
case "$MODEL" in *:*) FULL="$MODEL" ;; *) FULL="$MODEL:latest" ;; esac
MINE=$(printf '%s' "$PS" | awk -v k="\"$FULL\"" '{ i = index($0, k); if (i) print substr($0, i) }')
[ -z "$MINE" ] || PS="$MINE"
SIZE=$(jnum "$PS" size); VRAM=$(jnum "$PS" size_vram)
SHARE=$(awk -v s="${SIZE:-0}" -v v="${VRAM:-}" 'BEGIN { if (s > 0 && v != "") printf "%.3f", v / s }')
GENS="" PPS=""
i=1
while [ "$i" -le "$RUNS" ]; do
  say "Run $i of $RUNS: generating $TOKENS tokens..."
  OUT=$(gen "$TOKENS") || die "Generation failed."
  EC=$(jnum "$OUT" eval_count); ED=$(jnum "$OUT" eval_duration)
  PC=$(jnum "$OUT" prompt_eval_count); PD=$(jnum "$OUT" prompt_eval_duration)
  T=$(awk -v c="${EC:-0}" -v d="${ED:-0}" 'BEGIN { if (d > 0) printf "%.2f", c / (d / 1e9) }')
  P=$(awk -v c="${PC:-0}" -v d="${PD:-0}" 'BEGIN { if (d > 0 && c > 8) printf "%.1f", c / (d / 1e9) }')
  [ -n "$T" ] && GENS="$GENS $T"; [ -n "$P" ] && PPS="$PPS $P"
  i=$((i + 1))
done
med() { printf '%s\n' $1 | sort -n | awk '{ a[NR] = $1 } END { if (NR) { if (NR % 2) print a[(NR + 1) / 2]; else printf "%.2f\n", (a[NR / 2] + a[NR / 2 + 1]) / 2 } }'; }
GEN_TPS=$(med "$GENS"); PROMPT_TPS=$(med "$PPS")
[ -n "$GEN_TPS" ] || die "Ollama didn't report timings; update Ollama and try again."

say ""
say "$MODEL ($QUANT) on $GPU: $GEN_TPS tokens/s generating${PROMPT_TPS:+, $PROMPT_TPS tokens/s reading the prompt}."
if [ -n "$SHARE" ] && awk -v s="$SHARE" 'BEGIN { exit !(s < 0.99) }'; then
  say "Note: only $(awk -v s="$SHARE" 'BEGIN { printf "%d", s * 100 }')% of the model is on the GPU, so part of it runs from system RAM."
fi

num_or_null() { if [ -n "$1" ]; then printf '%s' "$1"; else printf 'null'; fi; }
BODY="{\"gpu\":\"$GPU\",\"mem_gb\":$(num_or_null "$MEM"),\"gpu_count\":$COUNT,\"model\":\"$(clean "$MODEL")\",\"quant\":\"$(clean "$QUANT")\",\"gen_tps\":$GEN_TPS,\"prompt_tps\":$(num_or_null "$PROMPT_TPS"),\"ctx\":$CTX,\"gpu_share\":$(num_or_null "$SHARE"),\"runtime\":\"ollama\",\"runtime_version\":\"$(clean "$OLLAMA_VERSION")\",\"os\":\"$OS\",\"client\":\"bench.sh/1\"}"

[ "$SEND" = "no" ] && { printf '%s\n' "$BODY"; exit 0; }
if [ "$SEND" = "ask" ]; then
  say ""; say "This is everything that would be sent to $SITE:"; say "  $BODY"
  [ -n "$TTY" ] || { say "Run again with --yes to share it."; exit 0; }
  ask "Share it so others can see real speeds for this hardware? [Y/n] "
  case "$REPLY" in [nN]*) say "Not sent."; exit 0 ;; esac
fi
RES=$(curl -sS -X POST "$SITE/api/bench" -H 'Content-Type: application/json' -d "$BODY") || die "Couldn't reach $SITE."
MSG=$(jstr "$RES" message); ERR=$(jstr "$RES" error)
if [ -n "$MSG" ]; then say "$MSG"; else say "Not accepted: ${ERR:-$RES}"; exit 1; fi