#!/usr/bin/env bash
# Put a locally-hosted model on a GPU node, configured the way the fleet needs it.
#
# The whole reason this is a script and not two commands in a README:
#
#   MEASURED on tank — stock `ornith:9b` reported input_tokens=2050 for a
#   48000-word prompt, and answered as though nothing had been dropped.
#
# Ollama defaults to a ~2K context window whatever the model card says, and it
# truncates SILENTLY. An agent turn is exactly the workload that hits that
# default, and the symptom is not an error: it is a model that answers
# confidently about a prompt it never saw. So the fleet never uses the stock tag.
# It uses a derived one with `num_ctx` pinned, and that tag is what
# `images/agent-ornith` bakes into ANTHROPIC_MODEL.
#
# Sizing, measured on a 16 GB RTX 5060 Ti with a 5.6 GB model:
#
#   num_ctx  32768 ->  6.3 GB resident
#   num_ctx  65536 ->  7.2 GB
#   num_ctx 131072 ->  9.3 GB   <- the default here
#   num_ctx 262144 -> 13.6 GB   (fits, 100% GPU, but leaves little headroom)
#
# Usage:
#   scripts/fleet-model-setup.sh <ssh-host> [base-model] [fleet-tag] [num_ctx]
#   scripts/fleet-model-setup.sh architect
set -uo pipefail

HOST="${1:?usage: $0 <ssh-host> [base-model] [fleet-tag] [num_ctx]}"
BASE="${2:-ornith:9b}"
TAG="${3:-ornith-fleet:9b}"
CTX="${4:-131072}"

die() { printf 'ABORT  %s\n' "$*" >&2; exit 2; }
ok()  { printf 'OK     %s\n' "$*"; }

ssh -o BatchMode=yes -o ConnectTimeout=10 "$HOST" true 2>/dev/null || die "cannot ssh to $HOST"

# A GPU is not optional. This model on CPU is slow enough that a mission would
# time out rather than fail, which is the worse outcome.
ssh "$HOST" 'command -v nvidia-smi >/dev/null && nvidia-smi -L' \
  || die "$HOST has no NVIDIA GPU — a local backend belongs on a GPU node"

ssh "$HOST" 'command -v ollama >/dev/null' \
  || die "$HOST has no ollama (the fleet page's update button installs it)"

# Version gate, with the reason. architect ran 0.30.5 and could not pull this
# model AT ALL — "requires a newer version of Ollama" — which is at least a loud
# failure. The quiet one is older still: /v1/messages only exists from v0.14.
have=$(ssh "$HOST" 'ollama --version' | grep -oE '[0-9]+\.[0-9]+\.[0-9]+' | head -1)
[ -n "$have" ] || die "could not read ollama's version on $HOST"
printf 'ollama %s on %s\n' "$have" "$HOST"

ssh "$HOST" "ollama pull '$BASE'" >/dev/null 2>&1 \
  || die "could not pull $BASE on $HOST — is ollama new enough for it?"
ok "$BASE pulled"

# The derived tag. `ollama create` is idempotent, so re-running is safe.
ssh "$HOST" "printf 'FROM %s\nPARAMETER num_ctx %s\n' '$BASE' '$CTX' > /tmp/Modelfile.fleet \
             && ollama create '$TAG' -f /tmp/Modelfile.fleet" >/dev/null 2>&1 \
  || die "could not create $TAG on $HOST"
ok "$TAG created with num_ctx=$CTX"

# PROVE the window, do not assume it. This is the check the whole script exists
# for: a Modelfile that silently failed to apply looks exactly like one that
# worked, right up until an agent loses its context mid-mission.
words=$((CTX / 2))
reported=$(ssh "$HOST" "python3 - <<'PY'
import json, urllib.request
filler = 'the quick brown fox jumps over the lazy dog ' * ($words // 9)
body = {'model': '$TAG', 'max_tokens': 8,
        'messages': [{'role': 'user', 'content': 'Log:\n' + filler + '\nReply OK.'}]}
req = urllib.request.Request('http://127.0.0.1:11434/v1/messages',
    data=json.dumps(body).encode(),
    headers={'content-type': 'application/json', 'x-api-key': 'ollama'})
try:
    print(json.load(urllib.request.urlopen(req, timeout=1800))['usage']['input_tokens'])
except Exception as e:
    print('ERR', e)
PY")
case "$reported" in
  ''|*[!0-9]*) die "context probe failed on $HOST: $reported" ;;
esac
# Allow for tokenisation slack, but nothing like a truncation. 2050 is what the
# stock tag reports for ANY prompt above it; anything near that means the
# num_ctx did not take.
floor=$((words / 2))
[ "$reported" -gt "$floor" ] \
  || die "$TAG reported only $reported input tokens for a ~$words-word prompt — \
the context window did not take, and it would truncate silently in a mission"
ok "context proven: ~$words words -> $reported input tokens"

# Tool calling, the other thing a mission agent cannot work without.
tools=$(ssh "$HOST" "curl -s -m 300 http://127.0.0.1:11434/v1/messages \
  -H 'content-type: application/json' -H 'x-api-key: ollama' \
  -d '{\"model\":\"$TAG\",\"max_tokens\":300,\"tools\":[{\"name\":\"read_file\",\"description\":\"Read a file\",\"input_schema\":{\"type\":\"object\",\"properties\":{\"path\":{\"type\":\"string\"}},\"required\":[\"path\"]}}],\"messages\":[{\"role\":\"user\",\"content\":\"Read src/lib.rs using the tool.\"}]}' \
  | python3 -c \"import json,sys; d=json.load(sys.stdin); print(d.get('stop_reason'), [b.get('type') for b in d.get('content',[])])\"")
case "$tools" in
  tool_use*) ok "tool calling proven: $tools" ;;
  *) die "no tool_use from $TAG ($tools) — Claude Code cannot drive a model that cannot call tools" ;;
esac

printf '\n%s is ready on %s. Next: build the rootfs there —\n' "$TAG" "$HOST"
printf '  ssh %s "cd ~/clawmates && docker build -f images/agent-ornith/Dockerfile -t clawmates/agent-ornith:dev images/agent-ornith/"\n' "$HOST"
printf '  scripts/fc-build-rootfs.sh %s clawmates/agent-ornith:dev local-ornith 8G\n' "$HOST"