#!/usr/bin/env bash
# scripts/benchmark.sh — Benchmark Ollama model inference speed
#
# Measures tokens/sec for each installed model.
# Results written to benchmark-results.md
#
# Usage:
#   bash scripts/benchmark.sh
#   bash scripts/benchmark.sh qwen2.5:7b   # benchmark one model

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO_DIR="$(cd "${SCRIPT_DIR}/../.." && pwd)"
source "${REPO_DIR}/scripts/common.sh"

RESULTS_FILE="${SCRIPT_DIR}/../benchmark-results.md"
PROMPT="Tell me a short story about a robot who loves cooking. Keep it to exactly 200 words."

if ! command_exists ollama; then
  die "Ollama not found."
fi

if ! curl -sf http://localhost:11434 -o /dev/null; then
  die "Ollama is not running."
fi

benchmark_model() {
  local model="$1"
  log_step "Benchmarking $model..."

  local start end elapsed
  start=$(date +%s%3N)

  local response
  response=$(ollama run "$model" "$PROMPT" 2>&1) || {
    log_error "Model $model failed to run."
    echo "| $model | ERROR | — |"
    return
  }

  end=$(date +%s%3N)
  elapsed=$(( (end - start) ))

  local word_count
  word_count=$(echo "$response" | wc -w)
  local tokens_est=$(( word_count * 4 / 3 ))  # rough estimate: 1 token ≈ 0.75 words
  local elapsed_sec
  elapsed_sec=$(echo "scale=1; $elapsed / 1000" | bc)
  local tps
  tps=$(echo "scale=1; $tokens_est / ($elapsed / 1000)" | bc 2>/dev/null || echo "?")

  printf "  %-30s  %6s tok/s  (%ss)\n" "$model" "$tps" "$elapsed_sec"
  echo "| \`$model\` | ${tps} tok/s | ${elapsed_sec}s |"
}

log_section "Ollama Benchmark"
log_info "Prompt: '$PROMPT'"
echo ""

if [[ -n "${1:-}" ]]; then
  models=("$@")
else
  # Get list of installed models
  mapfile -t models < <(ollama list 2>/dev/null | tail -n +2 | awk '{print $1}')
fi

if [[ ${#models[@]} -eq 0 ]]; then
  die "No models installed. Run: bash scripts/pull-models.sh"
fi

{
  echo "# Ollama Benchmark Results"
  echo "> Generated: $(date)"
  echo ""
  echo "| Model | Speed | Time for ~200 tok |"
  echo "|---|---|---|"
} > "$RESULTS_FILE"

for model in "${models[@]}"; do
  benchmark_model "$model" | tee -a "$RESULTS_FILE"
done

echo "" >> "$RESULTS_FILE"

log_success "Results written to $RESULTS_FILE"
echo ""
cat "$RESULTS_FILE"