#!/usr/bin/env bash
# Purpose: build llama.cpp from source with the GPU backend for one platform track, then
#          prove with a smoke test that the backend it was built for is the one in use
# Platform: spark (CUDA on aarch64), strix (Vulkan by default, HIP with BACKEND=hip),
#           mac (Metal, on by default), nvidia (CUDA on Linux or inside WSL2)
# Minimum memory: 8 GB
# Assumes: git, cmake and a C++ toolchain are installed; nvcc on PATH for the CUDA tracks;
#          glslc and the Vulkan headers for the Vulkan build, or hipconfig for the HIP
#          build; about 15 GB of free disk for the checkout and the build tree
#
# Usage: TRACK=<spark|strix|mac|nvidia> bash build-llama-cpp.sh
#        TRACK=strix BACKEND=hip bash build-llama-cpp.sh
#        TRACK=nvidia SMOKE_MODEL=~/models/unsloth/Qwen3-8B-GGUF/Qwen3-8B-Q4_K_M.gguf \
#            bash build-llama-cpp.sh
#
# Environment:
#   LLAMA_DIR    where the checkout lives              (default: $HOME/llama.cpp)
#   LLAMA_REF    tag, branch or commit to build        (default: master)
#   BUILD_DIR    the cmake build tree                  (default: $LLAMA_DIR/build)
#   JOBS         parallel compile jobs                 (default: detected)
#   GPU_TARGET   gfx target for the HIP build          (default: gfx1151)
#   SMOKE_MODEL  a GGUF file for the smoke test        (default: none; skipped)
#
# The cmake flags come from the project's build documentation, cited on the lesson page.
# Re-running is safe: the checkout is updated and cmake reconfigures in place.

set -euo pipefail

TRACK="${TRACK:-}"
BACKEND="${BACKEND:-}"
LLAMA_DIR="${LLAMA_DIR:-$HOME/llama.cpp}"
LLAMA_REF="${LLAMA_REF:-master}"
BUILD_DIR="${BUILD_DIR:-$LLAMA_DIR/build}"
GPU_TARGET="${GPU_TARGET:-gfx1151}"
SMOKE_MODEL="${SMOKE_MODEL:-}"

die() { echo "build-llama-cpp: $*" >&2; exit 1; }
have() { command -v "$1" >/dev/null 2>&1; }

[ -n "$TRACK" ] || die "set TRACK=spark, TRACK=strix, TRACK=mac or TRACK=nvidia"
have git || die "git is not installed"
have cmake || die "cmake is not installed"

if [ -z "${JOBS:-}" ]; then
  if have nproc; then
    JOBS="$(nproc)"
  elif have sysctl; then
    JOBS="$(sysctl -n hw.ncpu)"
  else
    JOBS=4
  fi
fi

# --- 1. Decide the backend and the cmake flags for this track ------------------------
CMAKE_FLAGS=(-DCMAKE_BUILD_TYPE=Release)
case "$TRACK" in
  spark|nvidia)
    have nvcc || die "nvcc is not on PATH; install the CUDA toolkit before building the CUDA backend"
    CMAKE_FLAGS+=(-DGGML_CUDA=ON)
    EXPECT_BACKEND="CUDA"
    ;;
  strix)
    if [ "$BACKEND" = "hip" ]; then
      have hipconfig || die "hipconfig is not on PATH; install ROCm, or leave BACKEND unset to build Vulkan"
      HIPCXX="$(hipconfig -l)/clang"
      HIP_PATH="$(hipconfig -R)"
      export HIPCXX HIP_PATH
      CMAKE_FLAGS+=(-DGGML_HIP=ON "-DGPU_TARGETS=$GPU_TARGET")
      EXPECT_BACKEND="ROCm"
    else
      have glslc || die "glslc is not on PATH; install the Vulkan development packages (see the lesson page)"
      CMAKE_FLAGS+=(-DGGML_VULKAN=1)
      EXPECT_BACKEND="Vulkan"
    fi
    ;;
  mac)
    [ "$(uname -s)" = "Darwin" ] || die "TRACK=mac only makes sense on macOS; this machine reports $(uname -s)"
    # Metal is enabled by default on macOS, so no extra flag is needed here.
    EXPECT_BACKEND="Metal"
    ;;
  *)
    die "unknown TRACK '$TRACK'; use spark, strix, mac or nvidia"
    ;;
esac

echo "==> Track $TRACK, backend $EXPECT_BACKEND, ${JOBS} parallel jobs"
echo "    cmake flags: ${CMAKE_FLAGS[*]}"

# --- 2. Get or update the source ------------------------------------------------------
if [ -d "$LLAMA_DIR/.git" ]; then
  echo "==> Updating $LLAMA_DIR"
  git -C "$LLAMA_DIR" fetch --tags origin
else
  echo "==> Cloning llama.cpp into $LLAMA_DIR"
  git clone https://github.com/ggml-org/llama.cpp "$LLAMA_DIR"
fi
git -C "$LLAMA_DIR" checkout "$LLAMA_REF"
git -C "$LLAMA_DIR" pull --ff-only 2>/dev/null || true
BUILT_REF="$(git -C "$LLAMA_DIR" rev-parse --short HEAD)"
echo "    building $LLAMA_REF at $BUILT_REF"

# --- 3. Configure and build -----------------------------------------------------------
echo "==> Configuring"
cmake -S "$LLAMA_DIR" -B "$BUILD_DIR" "${CMAKE_FLAGS[@]}"
echo "==> Building (this takes several minutes the first time)"
cmake --build "$BUILD_DIR" --config Release -j "$JOBS"

BIN="$BUILD_DIR/bin"
[ -x "$BIN/llama-cli" ] || die "the build finished but $BIN/llama-cli is missing; read the build output above"
echo "==> Binaries in $BIN:"
find "$BIN" -maxdepth 1 -type f -perm -u+x -exec basename {} \; | sort | sed 's/^/    /'

# --- 4. Smoke test: prove which backend is in use --------------------------------------
echo "==> Version and build information"
"$BIN/llama-cli" --version 2>&1 | sed 's/^/    /'

if [ -z "$SMOKE_MODEL" ]; then
  cat <<'MSG'

No SMOKE_MODEL was given, so the backend has not been exercised yet. Download a small
GGUF file and re-run with SMOKE_MODEL=<path> to see the backend column, or run
llama-bench yourself as the lesson page shows.
MSG
  exit 0
fi

[ -f "$SMOKE_MODEL" ] || die "SMOKE_MODEL '$SMOKE_MODEL' does not exist"
echo "==> Smoke test: eight generated tokens, one repetition"
SMOKE_OUT="$("$BIN/llama-bench" -m "$SMOKE_MODEL" -p 0 -n 8 -r 1 2>&1 || true)"
printf '%s\n' "$SMOKE_OUT" | sed 's/^/    /'

if printf '%s' "$SMOKE_OUT" | grep -qi "$EXPECT_BACKEND"; then
  echo "==> The backend column names $EXPECT_BACKEND. The build is using the GPU path it was built for."
else
  echo "==> WARNING: '$EXPECT_BACKEND' does not appear in the llama-bench output above." >&2
  echo "    The build may have fallen back to the CPU. See the challenge page for the procedure." >&2
  exit 1
fi