#!/usr/bin/env bash
# =============================================================================
# Gemma-4-12B-it One-Shot Setup for Pascal GPUs (GTX 1070, 8GB VRAM)
# Linux (Ubuntu 24.04 / Debian-based)
#
# Usage: chmod +x setup_gemma4_pascal.sh && ./setup_gemma4_pascal.sh
#
# This does everything: installs deps, builds llama.cpp, downloads the model,
# and runs a test. Add --no-download to skip model download when testing.
# =============================================================================
set -e

# ─── Config ──────────────────────────────────────────────────────────────────
INSTALL_DIR="${HOME}/llama.cpp"
MODELS_DIR="${HOME}/models"
MODEL_URL="https://huggingface.co/lmstudio-community/gemma-4-12B-it-GGUF/resolve/main/gemma-4-12B-it-Q4_K_M.gguf"
MMPROJ_URL="https://huggingface.co/lmstudio-community/gemma-4-12B-it-GGUF/resolve/main/mmproj-gemma-4-12B-it-BF16.gguf"

# ─── Color output ────────────────────────────────────────────────────────────
RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'; NC='\033[0m'
info()  { echo -e "${GREEN}[✓]${NC} $1"; }
warn()  { echo -e "${YELLOW}[!]${NC} $1"; }
err()   { echo -e "${RED}[✗]${NC} $1"; }
step()  { echo -e "\n${YELLOW}══ $1 ══${NC}"; }

# ─── 1. Install system deps ─────────────────────────────────────────────────
step "1/6 — Installing system dependencies"
if command -v apt-get &>/dev/null; then
  sudo DEBIAN_FRONTEND=noninteractive apt-get update -qq
  sudo DEBIAN_FRONTEND=noninteractive apt-get install -y -qq git cmake build-essential curl psmisc
elif command -v dnf &>/dev/null; then
  sudo dnf install -y git cmake make gcc-c++ curl psmisc
elif command -v pacman &>/dev/null; then
  sudo pacman -S --noconfirm git cmake base-devel curl psmisc
else
  warn "Unknown package manager — ensure git, cmake, make, g++, curl are installed"
fi
info "System dependencies ready"

# ─── 2. Check / install CUDA ────────────────────────────────────────────────
step "2/6 — Checking CUDA toolkit"
if ! command -v nvcc &>/dev/null; then
  warn "nvcc not found in PATH — checking conda..."
  if [ -d "$HOME/miniconda" ] && [ -f "$HOME/miniconda/bin/nvcc" ]; then
    export PATH="$HOME/miniconda/bin:$PATH"
    export CUDA_HOME="$HOME/miniconda"
    export CMAKE_CUDA_COMPILER="$HOME/miniconda/bin/nvcc"
    info "Found nvcc via conda: $(nvcc --version | tail -1)"
  elif [ -d "$HOME/miniconda3" ] && [ -f "$HOME/miniconda3/bin/nvcc" ]; then
    export PATH="$HOME/miniconda3/bin:$PATH"
    export CUDA_HOME="$HOME/miniconda3"
    export CMAKE_CUDA_COMPILER="$HOME/miniconda3/bin/nvcc"
    info "Found nvcc via conda3: $(nvcc --version | tail -1)"
  else
    err "CUDA toolkit not found!"
    echo "  Install from: https://developer.nvidia.com/cuda-downloads"
    echo "  Or via conda: conda install -c conda-forge cuda-toolkit"
    exit 1
  fi
else
  info "nvcc found: $(nvcc --version | tail -1)"
fi
export LD_LIBRARY_PATH="${CUDA_HOME:-/usr/local/cuda}/lib64:${LD_LIBRARY_PATH:-}"

# ─── 3. Clone and build llama.cpp ───────────────────────────────────────────
step "3/6 — Building llama.cpp from source (this takes 3-10 minutes)"
if [ -d "$INSTALL_DIR" ]; then
  warn "Removing existing build at $INSTALL_DIR"
  rm -rf "$INSTALL_DIR"
fi

git clone --depth 1 https://github.com/ggml-org/llama.cpp.git "$INSTALL_DIR"
cd "$INSTALL_DIR"

# -DCMAKE_CUDA_ARCHITECTURES="61" is REQUIRED on Pascal (CC 6.1). Without it
# cmake auto-detect targets CC 7.0+ and the binary silently runs on the CPU.
cmake -B build \
  -DGGML_CUDA=ON \
  -DCMAKE_BUILD_TYPE=Release \
  -DCMAKE_CUDA_ARCHITECTURES="61" \
  ${CUDA_HOME:+-DCUDA_HOME="$CUDA_HOME"} \
  ${CMAKE_CUDA_COMPILER:+-DCMAKE_CUDA_COMPILER="$CMAKE_CUDA_COMPILER"}

cmake --build build -j$(nproc) --target llama-cli llama-mtmd-cli llama-bench 2>&1 | \
  awk '{if(NR%20==0) printf "."} END{print " done"}'

CLI="$INSTALL_DIR/build/bin/llama-cli"
MTMD="$INSTALL_DIR/build/bin/llama-mtmd-cli"

if [ ! -f "$CLI" ] || [ ! -f "$MTMD" ]; then
  err "Build failed — binaries not found"
  exit 1
fi

VERSION=$($CLI --version 2>&1 | grep version || echo "unknown")
info "llama.cpp built successfully: $VERSION"

# ─── 4. Download model ──────────────────────────────────────────────────────
step "4/6 — Downloading Gemma-4-12B-it Q4_K_M (6.86 GB + 168 MB mmproj)"
mkdir -p "$MODELS_DIR"

MODEL_FILE="$MODELS_DIR/gemma-4-12B-it-Q4_K_M.gguf"
MMPROJ_FILE="$MODELS_DIR/mmproj-gemma-4-12B-it-BF16.gguf"

if [ ! -f "$MODEL_FILE" ]; then
  echo "  Downloading model (6.86 GB)..."
  # -C - resumes a partial download instead of restarting from zero if the
  # 6.86 GB transfer is interrupted.
  curl -L -C - -o "$MODEL_FILE" "$MODEL_URL"
  info "Model downloaded: $(du -h "$MODEL_FILE" | cut -f1)"
else
  warn "Model already exists, skipping download ($(du -h "$MODEL_FILE" | cut -f1))"
fi

if [ ! -f "$MMPROJ_FILE" ]; then
  echo "  Downloading mmproj (168 MB)..."
  curl -L -C - -o "$MMPROJ_FILE" "$MMPROJ_URL"
  info "mmproj downloaded: $(du -h "$MMPROJ_FILE" | cut -f1)"
else
  warn "mmproj already exists, skipping download"
fi

# Verify files
file "$MODEL_FILE" | grep -q GGUF && info "Model file is valid GGUF" || warn "Model file might be incomplete"
file "$MMPROJ_FILE" | grep -q data && info "mmproj file looks valid" || warn "mmproj might be incomplete"

# ─── 5. Create helper alias ─────────────────────────────────────────────────
step "5/6 — Creating helper scripts"
cat > "$MODELS_DIR/gemma4-chat.sh" << 'SCRIPT'
#!/bin/bash
export LD_LIBRARY_PATH="${HOME}/miniconda/lib:${LD_LIBRARY_PATH:-}"
MTMD="${HOME}/llama.cpp/build/bin/llama-mtmd-cli"

if [ $# -eq 0 ]; then
  # Interactive chat mode (no args)
  $MTMD --mmproj "${HOME}/models/mmproj-gemma-4-12B-it-BF16.gguf" \
    -m "${HOME}/models/gemma-4-12B-it-Q4_K_M.gguf" \
    -ngl 99 -fa on --no-kv-offload --jinja \
    -n 256
elif [ -f "$1" ]; then
  # First arg is an image file — analyze it. (Don't name the var `shift`: it
  # shadows the bash builtin used on the next line.)
  img_file="$1"; shift
  $MTMD --mmproj "${HOME}/models/mmproj-gemma-4-12B-it-BF16.gguf" \
    -m "${HOME}/models/gemma-4-12B-it-Q4_K_M.gguf" \
    -ngl 99 -fa on --no-kv-offload --jinja \
    --image "$img_file" \
    -p "$*" -n 300
else
  # Text prompt
  CLI="${HOME}/llama.cpp/build/bin/llama-cli"
  $CLI -m "${HOME}/models/gemma-4-12B-it-Q4_K_M.gguf" \
    -ngl 99 -fa on --no-kv-offload --jinja \
    -p "$*" -n 200 -e --no-display-prompt
fi
SCRIPT
chmod +x "$MODELS_DIR/gemma4-chat.sh"

# Add to .bashrc if not already there
if ! grep -q "gemma4-chat" "$HOME/.bashrc" 2>/dev/null; then
  echo "" >> "$HOME/.bashrc"
  echo "# Gemma-4-12B-it helper" >> "$HOME/.bashrc"
  echo 'alias gemma4="'"$MODELS_DIR"'/gemma4-chat.sh"' >> "$HOME/.bashrc"
  info "Added 'gemma4' alias to ~/.bashrc (run 'source ~/.bashrc' or open new terminal)"
fi

# ─── 6. Test ─────────────────────────────────────────────────────────────────
step "6/6 — Running test"
echo ""
echo "  ┌─────────────────────────────────────┐"
echo "  │  Testing Gemma-4-12B-it on Pascal   │"
echo "  └─────────────────────────────────────┘"
echo ""

kill $(fuser /dev/nvidia* 2>/dev/null | tr " " "\n" | sort -u) 2>/dev/null; sleep 1

export LD_LIBRARY_PATH="${CUDA_HOME:-/usr/local/cuda}/lib64:$HOME/miniconda/lib:${LD_LIBRARY_PATH:-}"

# Text test
echo "--- Text test ---"
$CLI -m "$MODEL_FILE" -ngl 99 -fa on --no-kv-offload --jinja \
  -p "Hello! What is 2+2?" -n 30 \
  -e 2>&1 | grep -v "^$\|^  /\|^> \|^|\|^\[.*t/s\]\|^\[.*ms\]\|^Exiting\|^Usage\|WARN\|^0\." | head -5

echo ""
echo "--- Done! ---"
echo ""
echo "  Setup complete! Commands:"
echo ""
echo "    gemma4 \"Your question\"           # Text prompt"
echo "    gemma4 photo.jpg \"Describe it\"    # Analyze an image"
echo "    gemma4                            # Interactive chat mode"
echo "    ~/models/gemma4-chat.sh            # Same, full path"
echo ""
echo "  Run 'source ~/.bashrc' then try: gemma4 \"Hello!\""
echo ""