#!/usr/bin/env bash
set -euo pipefail
IFS=$'\n\t'
trap 'echo "[ERROR] Line ${LINENO}: command failed" >&2' ERR

echo "[STEP] Updating APT & installing curl"
sudo apt update
sudo apt install -y curl

echo "[STEP] Checking for NVIDIA driver (nvidia-smi)"
if ! command -v nvidia-smi >/dev/null 2>&1; then
  echo "[WARN] nvidia-smi not found. Make sure you ran 01-nvidia_driver_install.sh and rebooted." >&2
else
  nvidia-smi || true
fi

echo "[STEP] Installing Ollama (official installer)"
if ! command -v ollama >/dev/null 2>&1; then
  # Official Ollama install script – sets up systemd service etc.
  curl -fsSL https://ollama.com/install.sh | sh
else
  echo "[INFO] Ollama already installed, skipping installer."
fi

echo "[STEP] Ensuring Ollama service is enabled & running"
if command -v systemctl >/dev/null 2>&1; then
  sudo systemctl enable --now ollama || true
  sudo systemctl status ollama --no-pager -l || true
else
  echo "[WARN] systemctl not available; make sure ollama daemon is running."
fi

echo
echo "[STEP] Pulling some example models (you can edit this list)"
MODELS=(
  "llama3"        # general chat / reasoning
  "codellama:7b"  # coding
  "qwen2-math:7b" # math reasoning
)

for model in "${MODELS[@]}"; do
  echo "[MODEL] Pulling: $model"
  ollama pull "$model"
done

echo
echo "[DONE] Ollama install finished."
echo "[TEST] Try:  ollama run llama3"
echo "[TIP] To expose Ollama on the network, you can run (for example):"
echo "      OLLAMA_HOST=0.0.0.0:11434 ollama serve"