#!/usr/bin/env bash
# Manage only the named experiment container. Never stop a production model.
set -Eeuo pipefail
CASE=${1:-}
ACTION=${2:-}
NAME=mike-ai-bonsai2-ab
IMAGE=mike-ai/bonsai2-ab:prism-b10685
MODEL=/data/models/bonsai2-ab/Ternary-Bonsai-2-27B-PQ2_0.gguf
GPU_5080=GPU-8ad38c6c-5a01-9d8e-1dfa-ed662ad78fbe
GPU_3060=GPU-4834d9d7-5b61-3004-1fb3-4ae49d482d4b
[[ $CASE =~ ^(fast|medium|ultra)$ && $ACTION =~ ^(create|start|stop)$ ]] || {
  echo 'Usage: case.sh {fast|medium|ultra} {create|start|stop}' >&2; exit 2;
}
test "$(hostname)" = athena || exit 2
if [[ $ACTION == stop ]]; then
  docker stop "$NAME" >/dev/null 2>&1 || true
  exit 0
fi
if [[ $ACTION == start ]]; then
  [[ ${3:-} == --go ]] || { echo 'Start blocked: explicit --go required' >&2; exit 2; }
  [[ -z $(docker ps --format '{{.Names}}' | grep -E '^mike-ai-llama-(fast|medium|large|ultra|uncensored)$' || true) ]] || {
    echo 'Refusing to start while a production LLM is running' >&2; exit 1;
  }
  [[ $(docker inspect -f '{{.State.Status}}' "$NAME") =~ ^(created|exited)$ ]]
  docker start "$NAME" >/dev/null
  exit 0
fi
test -s "$MODEL" || { echo 'Verified model is missing' >&2; exit 1; }
if docker inspect "$NAME" >/dev/null 2>&1; then
  test "$(docker inspect -f '{{.State.Running}}' "$NAME")" = false || {
    echo 'Refusing to replace a running experiment container' >&2; exit 1;
  }
  docker rm "$NAME" >/dev/null
fi
case "$CASE" in
  fast) CONTEXT=76800; UBATCH=64; DEVICES="$GPU_5080"; SPLIT=(--device CUDA0 --split-mode none) ;;
  medium) CONTEXT=160000; UBATCH=512; DEVICES="$GPU_5080,$GPU_3060"; SPLIT=(--device CUDA0,CUDA1 --main-gpu 0 --split-mode layer --tensor-split 85,15) ;;
  ultra) CONTEXT=262144; UBATCH=128; DEVICES="$GPU_5080,$GPU_3060"; SPLIT=(--device CUDA0,CUDA1 --main-gpu 0 --split-mode layer --tensor-split 80,20) ;;
esac
docker create --name "$NAME" --gpus "\"device=$DEVICES\"" \
  --network bridge -p 127.0.0.1:5006:8080 \
  --read-only --tmpfs /tmp:rw,noexec,nosuid,nodev,size=256m \
  --security-opt no-new-privileges:true --cap-drop ALL --pids-limit 1024 \
  --log-opt max-size=20m --log-opt max-file=2 \
  --label com.mike-ai.experiment=bonsai2-ab --label com.mike-ai.case="$CASE" \
  -e NVIDIA_VISIBLE_DEVICES="$DEVICES" -e NVIDIA_DRIVER_CAPABILITIES=compute,utility \
  -v /data/models/bonsai2-ab:/models:ro "$IMAGE" \
  --model "/models/$(basename "$MODEL")" --alias bonsai2-test \
  --ctx-size "$CONTEXT" --flash-attn on --cache-type-k q4_0 --cache-type-v q4_0 \
  --cache-prompt --threads 6 --threads-batch 6 --batch-size 2048 \
  --ubatch-size "$UBATCH" --parallel 1 --jinja --reasoning auto \
  --host 0.0.0.0 --port 8080 --metrics --fit off --n-gpu-layers all \
  --temperature 1.0 --top-p 0.95 --top-k 20 "${SPLIT[@]}" >/dev/null
echo "Created stopped Bonsai $CASE case ($CONTEXT context)."