#!/usr/bin/env bash
set -euo pipefail

LLAMA_SRC_DIR="${LLAMA_SRC_DIR:-/data/llama.cpp}"
LLAMA_BIN_DIR="${LLAMA_BIN_DIR:-/data/bin}"
LLAMA_DIFFUSION_BIN="${LLAMA_DIFFUSION_BIN:-${LLAMA_BIN_DIR}/llama-diffusion-cli}"
LLAMA_BUILD_CUDA="${LLAMA_BUILD_CUDA:-1}"
LLAMA_CMAKE_EXTRA_ARGS="${LLAMA_CMAKE_EXTRA_ARGS:-}"
LLAMA_CMAKE_EXE_LINKER_FLAGS="${LLAMA_CMAKE_EXE_LINKER_FLAGS:-}"
LLAMA_CMAKE_SHARED_LINKER_FLAGS="${LLAMA_CMAKE_SHARED_LINKER_FLAGS:-}"
CMAKE_BUILD_PARALLEL="${CMAKE_BUILD_PARALLEL:-$(nproc)}"

mkdir -p "${LLAMA_BIN_DIR}"

if [[ -x "${LLAMA_DIFFUSION_BIN}" ]]; then
  echo "[build] Existing llama-diffusion-cli: ${LLAMA_DIFFUSION_BIN}"
  exit 0
fi

if [[ ! -d "${LLAMA_SRC_DIR}/.git" ]]; then
  echo "[build] Cloning llama.cpp into ${LLAMA_SRC_DIR}"
  rm -rf "${LLAMA_SRC_DIR}"
  git clone --depth 1 https://github.com/ggml-org/llama.cpp "${LLAMA_SRC_DIR}"
fi

cd "${LLAMA_SRC_DIR}"

echo "[build] Checking out DiffusionGemma PR #24423"
git fetch origin pull/24423/head:diffusiongemma
git checkout diffusiongemma

CMAKE_ARGS=()
if [[ "${LLAMA_BUILD_CUDA}" == "1" || "${LLAMA_BUILD_CUDA}" == "true" ]]; then
  CMAKE_ARGS+=("-DGGML_CUDA=ON")
fi

NVIDIA_LIB_DIRS=()
while IFS= read -r lib_dir; do
  [[ -n "${lib_dir}" && -d "${lib_dir}" ]] && NVIDIA_LIB_DIRS+=("${lib_dir}")
done < <(python - <<'PY'
from pathlib import Path
import site

seen = set()
for root in [*site.getsitepackages(), site.getusersitepackages()]:
    if not root:
        continue
    base = Path(root) / "nvidia"
    if not base.exists():
        continue
    for path in base.glob("*/lib"):
        if path.is_dir():
            value = str(path)
            if value not in seen:
                seen.add(value)
                print(value)
PY
)

if [[ "${#NVIDIA_LIB_DIRS[@]}" -gt 0 ]]; then
  NVIDIA_LD_PATH="$(IFS=:; echo "${NVIDIA_LIB_DIRS[*]}")"
  export LD_LIBRARY_PATH="${NVIDIA_LD_PATH}${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}"

  LINK_FLAGS=""
  for lib_dir in "${NVIDIA_LIB_DIRS[@]}"; do
    LINK_FLAGS+=" -L${lib_dir} -Wl,-rpath,${lib_dir} -Wl,-rpath-link,${lib_dir}"
  done
  LINK_FLAGS+=" -Wl,--allow-shlib-undefined"
  CMAKE_ARGS+=("-DCMAKE_EXE_LINKER_FLAGS=${LINK_FLAGS} ${LLAMA_CMAKE_EXE_LINKER_FLAGS}")
  CMAKE_ARGS+=("-DCMAKE_SHARED_LINKER_FLAGS=${LINK_FLAGS} ${LLAMA_CMAKE_SHARED_LINKER_FLAGS}")
fi

if [[ -n "${LLAMA_CMAKE_EXTRA_ARGS}" ]]; then
  # shellcheck disable=SC2206
  EXTRA_ARGS=( ${LLAMA_CMAKE_EXTRA_ARGS} )
  CMAKE_ARGS+=("${EXTRA_ARGS[@]}")
fi

echo "[build] cmake args: ${CMAKE_ARGS[*]}"
cmake -B build -DCMAKE_BUILD_TYPE=Release "${CMAKE_ARGS[@]}"
cmake --build build -j "${CMAKE_BUILD_PARALLEL}" --config Release --target llama-diffusion-cli

if [[ ! -x build/bin/llama-diffusion-cli ]]; then
  echo "[build] build/bin/llama-diffusion-cli was not produced" >&2
  exit 1
fi

shopt -s nullglob
cp -a build/bin/lib*.so* "${LLAMA_BIN_DIR}/"
cp build/bin/llama-diffusion-cli "${LLAMA_DIFFUSION_BIN}"
chmod +x "${LLAMA_DIFFUSION_BIN}"
echo "[build] Installed ${LLAMA_DIFFUSION_BIN}"