#!/usr/bin/env bash
# Optional native toolchain; normal npm builds verify the committed artifacts.
set -euo pipefail
cd "$(dirname "$0")/../.."
project_root="$PWD"
build_root="${AI_BUILD_DIR:-$project_root/.runtime/ai-rebuild}"
if [[ -e "$build_root" ]]; then
  echo "Use a new AI_BUILD_DIR; existing build work is preserved." >&2
  exit 1
fi
mkdir -p "$build_root"
export UV_CACHE_DIR="$build_root/uv-cache"
if [[ ! -d "$build_root/mlc-llm/.git" ]]; then
  git clone https://github.com/mlc-ai/mlc-llm.git "$build_root/mlc-llm"
fi
git -C "$build_root/mlc-llm" checkout 2f78caa4db0f90730a11ee3bb5cbd5f23bf67f9f
export MLC_LLM_SOURCE_DIR="$build_root/mlc-llm"
export TVM_SOURCE_DIR="$MLC_LLM_SOURCE_DIR/3rdparty/tvm"
# TVM is explicitly paired to the compiler wheel, not the MLC submodule tip.
git -C "$MLC_LLM_SOURCE_DIR" submodule update --init --recursive
git -C "$TVM_SOURCE_DIR" checkout 837cb9de1127b48ce48e4cefe09e83215b9d4ba7
git -C "$TVM_SOURCE_DIR" submodule update --init --recursive
git -C "$TVM_SOURCE_DIR" apply "$project_root/scripts/ai/tvm-browser-memory.patch"
uv venv --python 3.13 "$build_root/venv"
uv pip install --python "$build_root/venv/bin/python" --no-deps --find-links https://mlc.ai/wheels/ -r scripts/ai/compiler-requirements.txt
# Build the exact FFI submodule: newer PyPI wheels crash this compiler ABI.
uv pip install --python "$build_root/venv/bin/python" --no-deps "$TVM_SOURCE_DIR/3rdparty/tvm-ffi"
python_site=$("$build_root/venv/bin/python" -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')
(cd "$python_site" && patch -p1 < "$project_root/scripts/ai/compiler-compat.patch")
if [[ "$(uname -s)" == Darwin ]]; then
  # Some upstream nightly Mach-O wheels have invalid signatures. Ad-hoc sign
  # only this dedicated environment's native libraries; no system changes.
  while IFS= read -r -d '' native; do
    codesign --verify "$native" 2>/dev/null || codesign --force --sign - "$native"
  done < <(find "$python_site/tvm" "$python_site/mlc_llm" -name '*.dylib' -print0)
fi
if [[ ! -d "$build_root/emsdk/.git" ]]; then
  git clone https://github.com/emscripten-core/emsdk.git "$build_root/emsdk"
fi
"$build_root/emsdk/emsdk" install 6.0.13
"$build_root/emsdk/emsdk" activate 6.0.13
source "$build_root/emsdk/emsdk_env.sh"
export EMCC_CFLAGS="-sDEFAULT_TO_CXX=1 -ffile-prefix-map=$MLC_LLM_SOURCE_DIR=mlc-llm -fdebug-prefix-map=$MLC_LLM_SOURCE_DIR=mlc-llm"
make -B -C "$MLC_LLM_SOURCE_DIR/web" EMCC="em++ -ffile-prefix-map=$MLC_LLM_SOURCE_DIR=mlc-llm" dist/wasm/mlc_wasm_runtime.bc
make -B -C "$TVM_SOURCE_DIR/web" EMCC="em++ -ffile-prefix-map=$MLC_LLM_SOURCE_DIR=mlc-llm" dist/wasm/wasm_runtime.bc dist/wasm/tvmjs_support.bc dist/wasm/webgpu_runtime.bc
make -C "$TVM_SOURCE_DIR/web" EMCC=em++ dist/wasm/tvmjs_runtime.wasi.js src/tvmjs_runtime_wasi.js
export TVM_HOME="$TVM_SOURCE_DIR" MLC_LLM_HOME="$build_root/cache"
export AI_MODEL_LIBRARY="$build_root/model-batch1.wasm"
# npm run build downloads and verifies the pinned model inputs first.
"$build_root/venv/bin/python" -m mlc_llm compile .runtime/ai-cache/mlc-chat-config.json --device webgpu --overrides max_batch_size=1 -o "$AI_MODEL_LIBRARY"
node scripts/ai/package-model.mjs