#!/usr/bin/env bash
# Build llama-server with CUDA for Colab (T4=75, L4=89).
set -euo pipefail

LLAMA_DIR="${1:-/content/llama.cpp}"
CUDA_ARCH="${2:-75}"

export CUDA_HOME="${CUDA_HOME:-/usr/local/cuda}"
export PATH="${CUDA_HOME}/bin:${PATH}"
export LD_LIBRARY_PATH="${CUDA_HOME}/lib64:${LD_LIBRARY_PATH:-}"

echo "[build] CUDA_HOME=${CUDA_HOME} arch=${CUDA_ARCH}"

if [[ ! -d "${LLAMA_DIR}/.git" ]]; then
  git clone --depth 1 https://github.com/ggml-org/llama.cpp.git "${LLAMA_DIR}"
else
  git -C "${LLAMA_DIR}" fetch origin master
  git -C "${LLAMA_DIR}" reset --hard origin/master
fi

rm -rf "${LLAMA_DIR}/build"

cmake -S "${LLAMA_DIR}" -B "${LLAMA_DIR}/build" \
  -DGGML_CUDA=ON \
  -DCMAKE_CUDA_ARCHITECTURES="${CUDA_ARCH}" \
  -DCMAKE_BUILD_TYPE=Release \
  -DCMAKE_CUDA_COMPILER="${CUDA_HOME}/bin/nvcc"

cmake --build "${LLAMA_DIR}/build" --config Release -j "$(nproc)" --target llama-server

BIN="${LLAMA_DIR}/build/bin/llama-server"
if [[ ! -x "${BIN}" ]]; then
  echo "[build] ERROR: ${BIN} not found"
  exit 1
fi

echo "[build] OK: ${BIN}"
"${BIN}" --version 2>/dev/null || true