#!/usr/bin/env bash
# Build the Sortformer diarization engine into .diarize/sortformer/.
#
# WHAT THIS BUILDS, AND WHY SO LITTLE OF IT: openresearchtools/engine is a whole
# local-AI runtime (Rust CLI, in-process llama bridge, PDF/VLM modules, a cluster
# controller). We want exactly one thing from it — the native ggml implementation
# of NVIDIA's streaming Sortformer diarizer in `diarize/addons/overlay`. That code
# links against `ggml` ALONE: no llama, no whisper, no Rust. So this script stages
# just `ggml` + `tools/realtime` + our driver and builds those, which is why it
# needs no part of upstream's build system (whose scripts are PowerShell/CUDA and
# would otherwise be the blocker on Linux).
#
# The upstream commit is PINNED. This is a research repo with no releases, and an
# unpinned clone would silently change what produced every sidecar on disk.
#
# Usage:
#   scripts/build-sortformer.sh                # Vulkan (default)
#   SORTFORMER_BACKEND=cpu scripts/build-sortformer.sh
#   SORTFORMER_JOBS=2 scripts/build-sortformer.sh
#
# Requires: git, cmake, a C++17 compiler, and for the Vulkan build the Vulkan
# loader + headers and `glslc` (Arch: vulkan-headers shaderc).

set -euo pipefail

REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
OUT_DIR="${SORTFORMER_DIR:-$REPO_ROOT/.diarize/sortformer}"
WORK_DIR="${SORTFORMER_BUILD_DIR:-$REPO_ROOT/.diarize/.sortformer-build}"

# Pinned 2026-08-10. Bumping this changes model execution, so re-run the
# reproducibility check in .diarize/README.md before trusting a new pin.
ENGINE_REPO="https://github.com/openresearchtools/engine"
ENGINE_COMMIT="8bb4928c429b76754c7c7cfbdec0c27419bef684"

MODEL_URL="https://huggingface.co/openresearchtools/diar_streaming_sortformer_4spk-v2.1-gguf/resolve/main/diar_streaming_sortformer_4spk-v2.1.gguf"
MODEL_NAME="diar_streaming_sortformer_4spk-v2.1.gguf"

BACKEND="${SORTFORMER_BACKEND:-vulkan}"
JOBS="${SORTFORMER_JOBS:-$(nproc 2>/dev/null || echo 4)}"

say() { printf '\033[1m==>\033[0m %s\n' "$*"; }

for tool in git cmake; do
  command -v "$tool" >/dev/null || { echo "missing required tool: $tool" >&2; exit 1; }
done

case "$BACKEND" in
  vulkan)
    command -v glslc >/dev/null || {
      echo "glslc not found — the Vulkan build compiles shaders with it." >&2
      echo "Install shaderc, or build CPU-only with SORTFORMER_BACKEND=cpu." >&2
      exit 1
    }
    CMAKE_BACKEND_FLAGS=(-DGGML_VULKAN=ON)
    ;;
  cpu)
    CMAKE_BACKEND_FLAGS=()
    ;;
  *)
    echo "SORTFORMER_BACKEND must be 'vulkan' or 'cpu' (got '$BACKEND')" >&2
    exit 1
    ;;
esac

mkdir -p "$WORK_DIR"

# 1. Upstream source, at the pin.
SRC="$WORK_DIR/engine"
if [ -d "$SRC/.git" ] && [ "$(git -C "$SRC" rev-parse HEAD 2>/dev/null)" = "$ENGINE_COMMIT" ]; then
  say "engine source already at $ENGINE_COMMIT"
else
  say "fetching engine @ ${ENGINE_COMMIT:0:12}"
  rm -rf "$SRC"
  git init -q "$SRC"
  git -C "$SRC" remote add origin "$ENGINE_REPO"
  # Blobless partial clone: the full history is ~200 MB and we need one commit.
  git -C "$SRC" fetch -q --depth 1 --filter=blob:none origin "$ENGINE_COMMIT"
  git -C "$SRC" checkout -q FETCH_HEAD
fi

# 2. Stage the minimal tree. Layout matters: tools/realtime/CMakeLists.txt refers
#    to ../../ggml/include, so ggml must sit beside tools/, exactly as upstream.
#
#    ggml is staged ONCE per pin and then left alone. Re-copying it would bump
#    every source mtime and force a full rebuild of the Vulkan shaders (~10 min)
#    every time only our driver changed. `cp -u` on the small realtime tree keeps
#    edits cheap while still picking up a new pin.
STAGE="$WORK_DIR/stage"
mkdir -p "$STAGE/tools"
if [ -f "$STAGE/.ggml-commit" ] && [ "$(cat "$STAGE/.ggml-commit")" = "$ENGINE_COMMIT" ]; then
  say "ggml already staged for this pin"
else
  say "staging ggml"
  rm -rf "$STAGE/ggml" "$STAGE/build"
  cp -r "$SRC/third_party/llama.cpp/ggml" "$STAGE/ggml"
  echo "$ENGINE_COMMIT" > "$STAGE/.ggml-commit"
fi
say "staging tools/realtime + driver"
mkdir -p "$STAGE/tools/realtime"
cp -ru "$SRC/diarize/addons/overlay/llama.cpp/tools/realtime/." "$STAGE/tools/realtime/"
cp -u "$REPO_ROOT/scripts/sortformer/diarize-file.cpp" "$STAGE/tools/realtime/diarize-file.cpp"

cat > "$STAGE/CMakeLists.txt" <<'EOF'
cmake_minimum_required(VERSION 3.14)
project(sortformer-diar C CXX)
set(CMAKE_CXX_STANDARD 17)
set(CMAKE_POSITION_INDEPENDENT_CODE ON)
add_subdirectory(ggml)
add_subdirectory(tools/realtime)
EOF

# Upstream's CMakeLists builds their parity tool; we append our driver rather than
# replacing it, so the staged tree stays a faithful copy plus one target. Guarded
# because the staged tree now persists across runs.
if ! grep -q 'add_executable(diarize-file' "$STAGE/tools/realtime/CMakeLists.txt"; then
  cat >> "$STAGE/tools/realtime/CMakeLists.txt" <<'EOF'

add_executable(diarize-file diarize-file.cpp)
target_compile_features(diarize-file PRIVATE cxx_std_17)
target_link_libraries(diarize-file PRIVATE ${TARGET_LIB})
EOF
fi

# 3. Build.
say "building ($BACKEND, -j$JOBS) — the Vulkan shader compile is the slow part"
cmake -S "$STAGE" -B "$STAGE/build" \
  -DCMAKE_BUILD_TYPE=Release \
  -DGGML_NATIVE=ON \
  "${CMAKE_BACKEND_FLAGS[@]}" >/dev/null
cmake --build "$STAGE/build" -j"$JOBS" --target diarize-file

# 4. Install the binary beside the ggml shared objects it loads, and point its
#    RPATH at $ORIGIN so it runs without LD_LIBRARY_PATH from any cwd.
say "installing to $OUT_DIR"
mkdir -p "$OUT_DIR"
cp "$STAGE/build/tools/realtime/diarize-file" "$OUT_DIR/diarize-file"
find "$STAGE/build" -name 'libggml*.so*' -exec cp -a {} "$OUT_DIR/" \;
if command -v patchelf >/dev/null; then
  patchelf --set-rpath '$ORIGIN' "$OUT_DIR/diarize-file"
else
  echo "note: patchelf not found; scripts/diarize-sortformer.mjs sets LD_LIBRARY_PATH itself" >&2
fi

# 5. Model.
if [ -f "$OUT_DIR/$MODEL_NAME" ]; then
  say "model already present"
else
  say "downloading model (~471 MB)"
  curl -fSL --progress-bar -o "$OUT_DIR/$MODEL_NAME.part" "$MODEL_URL"
  mv "$OUT_DIR/$MODEL_NAME.part" "$OUT_DIR/$MODEL_NAME"
fi

say "done"
echo
echo "  binary: $OUT_DIR/diarize-file"
echo "  model:  $OUT_DIR/$MODEL_NAME"
echo
echo "Wire it up in settings.json:"
echo "    \"diarization\": {"
echo "      \"engine\": \"sortformer\","
echo "      \"sortformerBin\":   \"$OUT_DIR/diarize-file\","
echo "      \"sortformerModel\": \"$OUT_DIR/$MODEL_NAME\","
echo "      \"backend\": \"$BACKEND\""
echo "    }"