#!/usr/bin/env bash
# Preparation only: downloads, verifies, builds and creates a STOPPED container.
set -Eeuo pipefail
cd "$(dirname "$0")"
MODEL_DIR=/data/models/bonsai2-ab
FILE=Ternary-Bonsai-2-27B-PQ2_0.gguf
REVISION=6ed5e12bf84b7a63069882c91dd9e9218647d17b
SHA256=3907dc1658db1f78a9826bf8d5bcb8dc65db0d466388937af57f2294fae62ec1
PROJECTOR=Ternary-Bonsai-2-27B-mmproj-Q8_0.gguf
PROJECTOR_SHA256=6807ede61d570bb86ba34b756a0fa109edc33668604de867c6ea6d8f1d631903
IMAGE=mike-ai/bonsai2-ab:prism-b10685
ARCHIVE=prism-cuda.tar.gz
ARCHIVE_SHA256=4ec1572702fa3fd359653528fa5625fdd9c7ae02dedab7146da7773dae48cf2c
ARCHIVE_URL=https://github.com/PrismML-Eng/llama.cpp/releases/download/prism-b10685-7dffb15/llama-prism-b10685-7dffb15-bin-linux-cuda-12.8-x64.tar.gz

test "$(hostname)" = athena || { echo 'Run on Athena only' >&2; exit 2; }
install -d -m 0755 "$MODEL_DIR"
if ! (cd "$MODEL_DIR" && printf '%s  %s\n' "$SHA256" "$FILE" | sha256sum -c --status); then
  curl --fail --location --retry 5 --continue-at - \
    --output "$MODEL_DIR/$FILE.part" \
    "https://huggingface.co/prism-ml/Ternary-Bonsai-2-27B-gguf/resolve/$REVISION/$FILE"
  (cd "$MODEL_DIR" && printf '%s  %s\n' "$SHA256" "$FILE.part" | sha256sum -c)
  mv "$MODEL_DIR/$FILE.part" "$MODEL_DIR/$FILE"
fi
if ! (cd "$MODEL_DIR" && printf '%s  %s\n' "$PROJECTOR_SHA256" "$PROJECTOR" | sha256sum -c --status); then
  curl --fail --location --retry 5 --continue-at - \
    --output "$MODEL_DIR/$PROJECTOR.part" \
    "https://huggingface.co/prism-ml/Ternary-Bonsai-2-27B-gguf/resolve/$REVISION/$PROJECTOR"
  (cd "$MODEL_DIR" && printf '%s  %s\n' "$PROJECTOR_SHA256" "$PROJECTOR.part" | sha256sum -c)
  mv "$MODEL_DIR/$PROJECTOR.part" "$MODEL_DIR/$PROJECTOR"
fi
if ! printf '%s  %s\n' "$ARCHIVE_SHA256" "$ARCHIVE" | sha256sum -c --status; then
  curl --fail --location --retry 5 --continue-at - --output "$ARCHIVE.part" "$ARCHIVE_URL"
  printf '%s  %s\n' "$ARCHIVE_SHA256" "$ARCHIVE.part" | sha256sum -c
  mv "$ARCHIVE.part" "$ARCHIVE"
fi
docker build --label com.mike-ai.experiment=bonsai2-ab -t "$IMAGE" .
./case.sh fast create
test "$(docker inspect -f '{{.State.Status}}' mike-ai-bonsai2-ab)" = created
echo 'Bonsai 2 is prepared; test container is stopped. No inference was run.'