#!/usr/bin/env bash
set -Eeuo pipefail
CASE=${1:-}
ACTION=${2:-}
NAME=mike-ai-ornith15-ab
IMAGE=mike-ai/llama.cpp:local
MODEL=/data/models/ornith15-ab/Ornith-1.5-35B-Q4_K_M.gguf
GPU_5080=GPU-8ad38c6c-5a01-9d8e-1dfa-ed662ad78fbe
GPU_3060=GPU-4834d9d7-5b61-3004-1fb3-4ae49d482d4b
[[ $CASE =~ ^(fast|medium|ultra|vision)$ && $ACTION =~ ^(create|start|stop)$ ]] || {
  echo 'Usage: case.sh {fast|medium|ultra|vision} {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 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|vision) CONTEXT=76800; UBATCH=64; SPLIT=70,30 ;;
  medium) CONTEXT=160000; UBATCH=512; SPLIT=68,32 ;;
  ultra) CONTEXT=262144; UBATCH=128; SPLIT=65,35 ;;
esac
VISION=()
if [[ $CASE == vision ]]; then
  test -s /data/models/ornith15-ab/mmproj-Ornith-1.5-35B-BF16.gguf || {
    echo 'Verified vision projector is missing' >&2; exit 1;
  }
  VISION=(--mmproj /models/mmproj-Ornith-1.5-35B-BF16.gguf --mmproj-offload --mmproj-device CUDA1)
fi
docker create --name "$NAME" --gpus "\"device=$GPU_5080,$GPU_3060\"" \
  --ipc host --network bridge -p 127.0.0.1:5007:8080 \
  --read-only --tmpfs /tmp:size=1g,mode=1777 \
  --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=ornith15-ab --label com.mike-ai.case="$CASE" \
  -e NVIDIA_VISIBLE_DEVICES="$GPU_5080,$GPU_3060" -e NVIDIA_DRIVER_CAPABILITIES=compute,utility \
  -e MTMD_BACKEND_DEVICE=CUDA1 \
  -v /data/models/ornith15-ab:/models:ro "$IMAGE" \
  --model "/models/$(basename "$MODEL")" --alias ornith15-test \
  --ctx-size "$CONTEXT" --flash-attn on --cache-type-k q4_0 --cache-type-v q4_0 \
  --cache-prompt --cache-ram 32768 --threads 6 --threads-batch 6 --batch-size 2048 \
  --ubatch-size "$UBATCH" --parallel 1 --kv-unified --jinja --reasoning auto --reasoning-preserve \
  --reasoning-effort medium --reasoning-budget 2048 \
  --host 0.0.0.0 --port 8080 --metrics --fit off --n-gpu-layers all --load-mode none --no-ui \
  --temperature 0.2 --top-p 0.8 --top-k 20 \
  --device CUDA0,CUDA1 --main-gpu 0 --split-mode layer --tensor-split "$SPLIT" \
  --spec-type draft-mtp --spec-draft-n-max 3 --spec-draft-type-k f16 --spec-draft-type-v f16 \
  "${VISION[@]}" >/dev/null
echo "Created stopped Ornith $CASE case ($CONTEXT context)."