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

ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
STRATA_COMMIT="296b59d2f73c8f04a805df29c2fa97655adace78"
LLAMA_COMMIT="a0e6dc7384edb8f5d927155e1d26c10b4a310e6d"
METAL_CPP_URL="https://developer.apple.com/metal/cpp/files/metal-cpp_26.zip"
METAL_CPP_SHA256="4df3c078b9aadcb516212e9cb03004cbc5ce9a3e9c068fa3144d021db585a3a4"

mkdir -p "$ROOT/vendor"

if [ ! -d "$ROOT/vendor/llama.cpp/.git" ]; then
  git clone https://github.com/ggml-org/llama.cpp.git "$ROOT/vendor/llama.cpp"
fi
git -C "$ROOT/vendor/llama.cpp" fetch --all --tags
git -C "$ROOT/vendor/llama.cpp" checkout "$LLAMA_COMMIT"

if [ ! -d "$ROOT/vendor/strata/.git" ]; then
  git clone https://github.com/Niko1221/Strata.git "$ROOT/vendor/strata"
fi
git -C "$ROOT/vendor/strata" fetch --all --tags
git -C "$ROOT/vendor/strata" checkout -f "$STRATA_COMMIT"
git -C "$ROOT/vendor/strata" clean -fd
git -C "$ROOT/vendor/strata" apply "$ROOT/scripts/strata_metal_complete.diff"

if [ ! -f "$ROOT/vendor/strata/third_party/metal-cpp/Metal/Metal.hpp" ]; then
  TMP_ZIP="$(mktemp -t metal-cpp.XXXXXX.zip)"
  TMP_DIR="$(mktemp -d -t metal-cpp.XXXXXX)"
  curl -fL "$METAL_CPP_URL" -o "$TMP_ZIP"
  echo "$METAL_CPP_SHA256  $TMP_ZIP" | shasum -a 256 -c -
  unzip -q "$TMP_ZIP" -d "$TMP_DIR"
  mkdir -p "$ROOT/vendor/strata/third_party"
  rm -rf "$ROOT/vendor/strata/third_party/metal-cpp"
  mv "$TMP_DIR/metal-cpp" "$ROOT/vendor/strata/third_party/metal-cpp"
  rm -rf "$TMP_ZIP" "$TMP_DIR"
fi

cmake -S "$ROOT/vendor/strata/metal" -B "$ROOT/vendor/strata/build-metal" -G Ninja \
  -DCMAKE_BUILD_TYPE=Release \
  -DMETAL_CPP_DIR="$ROOT/vendor/strata/third_party/metal-cpp" \
  -DSTRATA_GGML_DIR="$ROOT/vendor/llama.cpp"

ninja -C "$ROOT/vendor/strata/build-metal" strata_window_cli convert_k256_q2_0 heal_k256_q2_0

cmake -S "$ROOT/vendor/strata/tools/vision" -B "$ROOT/vendor/strata/build-vision" -G Ninja \
  -DCMAKE_BUILD_TYPE=Release \
  -DLLAMA_DIR="$ROOT/vendor/llama.cpp" \
  -DSTRATA_VISION_CUDA=OFF \
  -DSTRATA_PORTABLE=OFF \
  -DGGML_METAL=ON

ninja -C "$ROOT/vendor/strata/build-vision" strata-vision

UV_BIN="${UV:-$(command -v uv || echo "$HOME/.local/bin/uv")}"
if [ "${SKIP_MODEL_DOWNLOAD:-0}" != "1" ]; then
  "$UV_BIN" run --with regex python3 "$ROOT/scripts/build_k512_iq2_xs_metal.py"
  "$UV_BIN" run --with numpy python3 "$ROOT/scripts/build_mtp_q2_0.py"
fi