#!/bin/sh
# The M3a check on straylight, run by the owner (not part of `make gate` or `verify-device`).
#
# A private home with one `ask` grant for `read_file` on a directory; `brokerd serve` and
# `loopd serve` on it, `loopd` talking to the real server through a private `inferproxy`. A
# `bxctl chat --say` asks Ornith to read a file in that directory; the approval appears in
# `bxctl approvals`; approving it gives the M3a runner's failure, which the model reports. The
# audit log must verify and hold a Decision, an Approval and a Result.
#
# It uses slot 0 only, and first checks that slot 0 is idle: the server is shared.
#
#   sh tools/check-m3a-device.sh [host:port]      (default straylight:11434)
#
# Needs curl and jq. On success the home is removed; on failure its path is printed. If the
# model is not loaded, `/slots` may fail: load it first (the check does not load models).
set -u

UPSTREAM="${1:-straylight:11434}"
MODEL="${BOXMAKER_MODEL:-ornith-1.5-35b-a3b}"
ROOT=$(cd "$(dirname "$0")/.." && pwd) || exit 1
BIN="$ROOT/target/debug"

fail() {
    echo "check-m3a-device: FAILED: $*" >&2
    [ -n "${HOME_DIR:-}" ] && echo "check-m3a-device: the home is kept at $HOME_DIR" >&2
    exit 1
}

for tool in curl jq cargo; do
    command -v "$tool" > /dev/null || fail "$tool is not installed"
done

# 1. Slot 0 must be idle. Anything but a clear "not processing" stops the check.
slots=$(curl -sf "http://$UPSTREAM/slots?model=$MODEL") || fail "cannot read /slots from $UPSTREAM"
# Not `jq -e`: it exits 1 when the value is `false`, which is the answer we want.
busy=$(printf '%s' "$slots" | jq '.[] | select(.id == 0) | .is_processing') \
    || fail "the /slots answer is not a list of slots"
[ -n "$busy" ] || fail "slot 0 is not in the /slots answer"
[ "$busy" = "false" ] || fail "slot 0 is busy ($busy); try again later"

# 2. Build.
cargo build --workspace --locked --manifest-path "$ROOT/Cargo.toml" || fail "cargo build"

HOME_DIR=$(mktemp -d) || fail "mktemp"
PIDS=""
cleanup() {
    for pid in $PIDS; do kill "$pid" 2> /dev/null; done
}
trap cleanup EXIT

wait_for() { # path, seconds
    n=0
    while [ ! -S "$1" ]; do
        n=$((n + 1))
        [ "$n" -gt $(($2 * 10)) ] && fail "$1 did not appear within $2 s"
        sleep 0.1
    done
}

# 3. The home: a file to read, one ask grant, the configs, the system prompt.
mkdir -p "$HOME_DIR/files" "$HOME_DIR/grants" "$HOME_DIR/run/infer" || fail "mkdir"
echo "The launch code is BANANA-42." > "$HOME_DIR/files/note.txt"
cat > "$HOME_DIR/grants/files-read.toml" <<EOF || fail "grant"
tool = "read_file"
mode = "ask"
max_taint = "secret"
result_class = "private"
untrusted = false

[constraints]
paths = ["$HOME_DIR/files"]
EOF
cat > "$HOME_DIR/brokerd.toml" <<EOF || fail "brokerd.toml"
[paths]
home = "$HOME_DIR"
grants = "$HOME_DIR/grants"
EOF
cat > "$HOME_DIR/config.toml" <<EOF || fail "config.toml"
[infer]
socket = "$HOME_DIR/run/infer/infer.sock"
model = "$MODEL"
[slots]
main = 0
background = 1
[expect]
template_sha256 = "f55f52930aa8bf44ab5cb85f99370fcc3c56e9a85640b812086d5330bce5d86b"
n_ctx = 131072
slots = 2
[paths]
home = "$HOME_DIR"
[broker]
socket = "$HOME_DIR/run/loop-broker/broker.sock"
EOF
cp "$ROOT/config/system.md" "$HOME_DIR/system.md" || fail "config/system.md"

# 4. The three daemons.
"$BIN/inferproxy" --listen "$HOME_DIR/run/infer/infer.sock" --upstream "$UPSTREAM" \
    2> "$HOME_DIR/inferproxy.err" &
PIDS="$PIDS $!"
wait_for "$HOME_DIR/run/infer/infer.sock" 5
"$BIN/brokerd" serve --config "$HOME_DIR/brokerd.toml" 2> "$HOME_DIR/brokerd.err" &
PIDS="$PIDS $!"
wait_for "$HOME_DIR/run/owner-broker/admin.sock" 10
"$BIN/loopd" serve --config "$HOME_DIR/config.toml" 2> "$HOME_DIR/loopd.err" &
PIDS="$PIDS $!"
wait_for "$HOME_DIR/run/loop/loop.sock" 60

ADMIN="$HOME_DIR/run/owner-broker/admin.sock"

# 5. The turn, in the background: it waits for the approval.
"$BIN/bxctl" chat --socket "$HOME_DIR/run/loop/loop.sock" --admin-socket "$ADMIN" \
    --session m3a-device --no-thinking \
    --say "Read the file $HOME_DIR/files/note.txt with the read_file tool and tell me exactly what happened." \
    > "$HOME_DIR/chat.out" 2> "$HOME_DIR/chat.err" &
CHAT=$!

# 6. Wait for the approval, check what it shows, approve it.
n=0
while :; do
    list=$("$BIN/bxctl" approvals --admin-socket "$ADMIN") || fail "bxctl approvals"
    [ "$list" != "no pending approvals" ] && break
    kill -0 "$CHAT" 2> /dev/null || fail "the turn ended without asking; see $HOME_DIR/chat.out"
    n=$((n + 1))
    [ "$n" -gt 300 ] && fail "no approval within 300 s"
    sleep 1
done
echo "$list"
printf '%s\n' "$list" | grep -q "grant files-read" || fail "the block does not name the grant"
printf '%s\n' "$list" | grep -q "read_file {\"path\":\"$HOME_DIR/files/note.txt\"}" \
    || fail "the block does not show the call"
id=$(printf '%s\n' "$list" | head -n 1 | cut -d ' ' -f 1)
"$BIN/bxctl" approve "$id" --admin-socket "$ADMIN" | tee "$HOME_DIR/approve.out"
grep -qx "approved $id: runs" "$HOME_DIR/approve.out" || fail "approve did not say it runs"

wait "$CHAT" || fail "bxctl chat failed; see $HOME_DIR/chat.err"
echo "--- the model's answer:"
cat "$HOME_DIR/chat.out"
echo "---"

# 7. The audit log verifies and holds the three records.
"$BIN/bxctl" audit verify --home "$HOME_DIR" || fail "the audit log does not verify"
for type in decision approval result; do
    cat "$HOME_DIR"/audit/*.jsonl | grep -q "\"type\":\"$type\"" || fail "no $type record"
done
grep -q "M3b" "$HOME_DIR/chat.out" \
    || echo "check-m3a-device: note: the model's answer does not quote the runner's sentence; read it above"

cleanup
trap - EXIT
rm -rf "$HOME_DIR"
echo "check-m3a-device: ok"