#!/bin/sh
set -eu

if [ "$#" -ne 3 ]; then
    echo "usage: $0 RUNNER STREAM_PROBE FAKE_SERVER" >&2
    exit 2
fi

runner=$1
stream_probe=$2
fake_server=$3
fixture=$(mktemp -d /private/tmp/gem-prefill-cancel-gate.XXXXXX)
cleanup()
{
    rm -rf "$fixture"
}
trap cleanup EXIT INT TERM HUP

mkdir -p "$fixture/build" "$fixture/scripts"
printf 'fixture model\n' > "$fixture/model.gguf"
printf 'fixture cancellation prompt\n' > "$fixture/prompt.txt"
# These fragments are the literal body of a generated executable fixture.
# shellcheck disable=SC2016
printf '%s\n' '#!/bin/sh' \
    'set -eu' \
    'if [ "$1" = print-config ]; then' \
    '    printf '\''{"engine":"gem-native-c-metal","context":512,"kv_type":"f32"}\n'\''' \
    '    exit 0' \
    'fi' \
    'port=' \
    'while [ "$#" -gt 0 ]; do' \
    '    if [ "$1" = --port ]; then port=$2; shift 2; else shift; fi' \
    'done' \
    'exec python3 "$GEM_FAKE_CANCEL_SERVER" --port "$port"' > "$fixture/build/gem"
printf '%s\n' '#!/bin/sh' 'printf '\''{"token_count":256}\n'\''' \
    > "$fixture/build/gem-tokenize"
chmod +x "$fixture/build/gem" "$fixture/build/gem-tokenize"
ln -s "$stream_probe" "$fixture/scripts/server-stream-probe.sh"

port=$((40000 + ($$ % 10000)))
(
    cd "$fixture"
    GEM_MODEL_SHA256=0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef \
        GEM_CANCEL_AFTER_SECONDS=1 GEM_CANCEL_RECOVERY_LIMIT_SECONDS=8 \
        GEM_FAKE_CANCEL_SERVER="$fake_server" \
        "$runner" model.gguf 512 prompt.txt results "$port" >/dev/null
)
jq -e '
  .workload == "prefill-cancel-recovery" and .context == 512 and
  .prompt_tokens == 256 and .cancellation_curl_status == 28 and
  .keepalive_events >= 1 and .post_cancel_health_status == 200 and
  .post_cancel_health_recovery_seconds <= 8 and
  .recovery_requested_max_tokens == 16 and
  .same_process_recovery_request == "passed"
' "$fixture/results/summary.json" >/dev/null
jq -e '
  .http_code == 200 and .requested_max_tokens == 16 and
  .visible_event_count > 0 and .finish_reason == "stop" and
  .usage == {prompt_tokens: 3, completion_tokens: 2, total_tokens: 5}
' "$fixture/results/recovery-measurement.json" >/dev/null
grep -q '^: keepalive$' "$fixture/results/cancelled.sse"