#!/bin/bash
# ollama-manager.sh
# Pull, list, benchmark, and prune Ollama models.
# Includes a size report so you know what is eating your disk.
#
# Usage: bash ollama-manager.sh [list|pull|bench|prune|size]

set -euo pipefail

OLLAMA_HOST="${OLLAMA_HOST:-http://localhost:11434}"

# ── Helpers ──────────────────────────────────────────────────────
check_ollama() {
  if ! curl -sf "${OLLAMA_HOST}" > /dev/null; then
    echo "ERROR: Ollama not reachable at ${OLLAMA_HOST}"
    echo "Is it running? Check: systemctl status ollama"
    exit 1
  fi
}

cmd_list() {
  check_ollama
  echo "Installed models:"
  echo ""
  curl -sf "${OLLAMA_HOST}/api/tags" | \
    python3 -c "
import json, sys
data = json.load(sys.stdin)
models = data.get('models', [])
if not models:
    print('  No models installed')
    sys.exit(0)
print(f'  {'NAME':<35} {'SIZE':>8}  {'MODIFIED':<20}')
print('  ' + '-'*70)
for m in models:
    name = m['name']
    size_gb = m['size'] / 1e9
    modified = m['modified_at'][:10]
    print(f'  {name:<35} {size_gb:>6.1f}GB  {modified:<20}')
print(f'\n  Total: {len(models)} model(s)')
"
}

cmd_pull() {
  check_ollama
  local model="${1:-}"
  if [[ -z "$model" ]]; then
    echo "Usage: $0 pull <model-name>"
    echo ""
    echo "Popular models:"
    echo "  llama3.2        3B params — fast, CPU-friendly"
    echo "  mistral         7B params — great balance"
    echo "  llama3.1:8b     8B params — strong reasoning"
    echo "  deepseek-coder-v2  coding specialist"
    echo "  phi3            3.8B Microsoft model — very efficient"
    exit 0
  fi
  echo "Pulling ${model}..."
  ollama pull "$model"
  echo "✓ Done"
}

cmd_bench() {
  check_ollama
  local model="${1:-llama3.2}"
  local prompt="Explain what a Linux process is in exactly one sentence."

  echo "Benchmarking: $model"
  echo "Prompt: \"$prompt\""
  echo ""

  start=$(date +%s%N)

  response=$(curl -sf "${OLLAMA_HOST}/api/generate" \
    -d "{\"model\":\"${model}\",\"prompt\":\"${prompt}\",\"stream\":false}" \
    2>/dev/null)

  end=$(date +%s%N)
  elapsed_ms=$(( (end - start) / 1000000 ))

  echo "$response" | python3 -c "
import json, sys
data = json.load(sys.stdin)
response = data.get('response', '')
tokens = data.get('eval_count', 0)
eval_ms = data.get('eval_duration', 1) / 1e6
tps = tokens / (eval_ms / 1000) if eval_ms > 0 else 0
print(f'Response: {response}')
print(f'')
print(f'Tokens generated: {tokens}')
print(f'Speed:            {tps:.1f} tokens/sec')
print(f'Total time:       {eval_ms:.0f}ms')
"
  echo "Wall time: ${elapsed_ms}ms"
}

cmd_size() {
  local models_dir="${HOME}/.ollama/models"
  if [[ ! -d "$models_dir" ]]; then
    models_dir="/usr/share/ollama/.ollama/models"
  fi

  echo "Ollama model storage:"
  echo ""
  if [[ -d "$models_dir" ]]; then
    du -sh "${models_dir}"/* 2>/dev/null | sort -rh | head -20 || \
      echo "  (no models found)"
    echo ""
    echo "Total:"
    du -sh "$models_dir" 2>/dev/null || echo "  unknown"
  else
    echo "Models directory not found: $models_dir"
  fi
}

cmd_prune() {
  check_ollama
  echo "Installed models (select one to remove):"
  echo ""

  mapfile -t models < <(
    curl -sf "${OLLAMA_HOST}/api/tags" | \
    python3 -c "
import json, sys
data = json.load(sys.stdin)
for m in data.get('models', []):
    size_gb = m['size'] / 1e9
    print(f\"{m['name']} ({size_gb:.1f}GB)\")
"
  )

  if [[ ${#models[@]} -eq 0 ]]; then
    echo "No models installed."
    exit 0
  fi

  for i in "${!models[@]}"; do
    echo "  $((i+1)). ${models[$i]}"
  done

  echo ""
  read -rp "Enter number to remove (or q to quit): " choice
  [[ "$choice" == "q" ]] && exit 0

  index=$((choice - 1))
  if [[ $index -lt 0 || $index -ge ${#models[@]} ]]; then
    echo "Invalid selection"
    exit 1
  fi

  model_name=$(echo "${models[$index]}" | awk '{print $1}')
  read -rp "Remove ${model_name}? [y/N] " confirm
  [[ "$confirm" =~ ^[Yy]$ ]] || exit 0

  ollama rm "$model_name"
  echo "✓ Removed ${model_name}"
}

# ── Main ─────────────────────────────────────────────────────────
case "${1:-list}" in
  list)  cmd_list ;;
  pull)  cmd_pull "${2:-}" ;;
  bench) cmd_bench "${2:-llama3.2}" ;;
  size)  cmd_size ;;
  prune) cmd_prune ;;
  *)
    echo "Usage: $0 [list|pull <model>|bench <model>|size|prune]"
    echo ""
    echo "  list          Show installed models with sizes"
    echo "  pull <model>  Download a model"
    echo "  bench <model> Benchmark a model (tokens/sec)"
    echo "  size          Show disk usage by model"
    echo "  prune         Interactively remove a model"
    ;;
esac