#!/bin/bash
# Stage and serve a local preview of the site (site/build.sh) with the
# cross-origin isolation headers, so WebAssembly threads work.
#
# Usage: dev-site.sh DIR              stage into DIR and serve on 127.0.0.1:$PORT
#        dev-site.sh DIR --rebuild    rebuild the page and app only
#                                     (the server keeps running)
#
# Required env:
#   MODEL     an export directory from train/export.sh (symlinked, not copied)
#   RUNTIME   the directory from train/vendor.sh
# Optional env:
#   OLDER     space-separated export directories staged as older brains the
#             app may keep running (to try the upgrade flow); newest first
#   PORT      default 8765
#   BASE      URL base path / directory for the site           (default: pet)
#   MODEL_ID  model directory name                  (default: basename of MODEL)
#   NEWS      the upgrade dialog's "what's new" line   (default: a test line)
#   FEATURES  JSON list of features the current brain has         (default: [])
#   TRAINED   text for the page's "Trained" line           (default: local build)
#   SOURCE    URL for the page's source link                       (default: #)
#   REMOTE_ENDPOINT  optional default Rust host URL                (default: empty)
#   ONNX_BYTES_PER_SECOND  throttle model downloads (default: 0, unlimited)
# The speed benchmark (test/bench.html) is served at /$BASE/bench.html.
# Nothing is published anywhere; this is for local testing.
set -euo pipefail
[ $# -ge 1 ] || { sed -n '2,/^set -e/p' "$0" | sed '$d;s/^# \{0,1\}//' >&2; exit 2; }
DIR=$(realpath -m "$1")
: "${MODEL:?set MODEL to an export directory}" "${RUNTIME:?set RUNTIME to the vendored runtime directory}"
HERE=$(cd "$(dirname "$0")" && pwd); REPO=$(dirname "$HERE")
PORT=${PORT:-8765} BASE=${BASE:-pet}
MODEL_ID=${MODEL_ID:-$(basename "$(realpath "$MODEL")")}
ROOT=$DIR/$BASE
mkdir -p "$ROOT"

# A release description for this preview (the real one is site/release.json).
SIZE=$(wc -c < "$MODEL/onnx/model_q4.onnx" | tr -d ' ')
python3 - "$DIR/release.json" "$MODEL_ID" "$SIZE" ${OLDER:-} <<'PY'
import json, os, sys
out, mid, size, *older = sys.argv[1:]
json.dump({"modelId": mid, "size": int(size), "trained": os.environ.get("TRAINED", "local build"),
           "news": os.environ.get("NEWS", "A preview brain, for testing the upgrade."),
           "features": json.loads(os.environ.get("FEATURES", "[]")),
           "older": [{"id": os.path.basename(os.path.realpath(o)), "features": []} for o in older],
           "runtime": "runtime", "source": os.environ.get("SOURCE", "#")}, open(out, "w"))
PY
RELEASE=$DIR/release.json BASE=/$BASE/ "$REPO/site/build.sh" "$ROOT"
[ "${2:-}" = "--rebuild" ] && exit 0

# Full previews refresh the local runtime exactly; published runtimes remain
# immutable. Finish a fresh copy before removing the previous preview runtime.
mkdir -p "$ROOT/runtime"
RUNTIME_STAGE=$(mktemp -d "$ROOT/runtime/.runtime.partial.XXXXXX")
trap 'rm -rf "$RUNTIME_STAGE"' EXIT
trap 'exit 1' HUP INT TERM
cp -a "$RUNTIME/." "$RUNTIME_STAGE/"
rm -rf "$ROOT/runtime/runtime"
mv "$RUNTIME_STAGE" "$ROOT/runtime/runtime"
trap - EXIT HUP INT TERM
for m in "$MODEL" ${OLDER:-}; do
  id=$(basename "$(realpath "$m")"); [ "$m" = "$MODEL" ] && id=$MODEL_ID
  mkdir -p "$ROOT/models/$id/onnx"
  cp "$m"/*.json "$m"/SHA256SUMS "$ROOT/models/$id/"
  ln -sf "$(realpath "$m")"/onnx/*.onnx "$ROOT/models/$id/onnx/"
done
echo "serving $DIR at http://127.0.0.1:$PORT/$BASE/"
exec python3 "$HERE/serve.py" "$DIR" "$PORT"