#!/bin/sh
set -eu

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

probe=$1
server=$2
port=$((40000 + $$ % 15000))
prompt=$(mktemp /private/tmp/gem-stream-probe-prompt.XXXXXX)
response=$(mktemp /private/tmp/gem-stream-probe-response.XXXXXX)
summary=$(mktemp /private/tmp/gem-stream-probe-summary.XXXXXX)
request=$(mktemp /private/tmp/gem-stream-probe-request.XXXXXX)
request_response=$(mktemp /private/tmp/gem-stream-probe-request-response.XXXXXX)
request_summary=$(mktemp /private/tmp/gem-stream-probe-request-summary.XXXXXX)
server_pid=
cleanup()
{
    if [ -n "$server_pid" ]; then
        kill "$server_pid" 2>/dev/null || true
        wait "$server_pid" 2>/dev/null || true
    fi
    rm -f "$prompt" "$response" "$summary" "$request" "$request_response" \
        "$request_summary"
}
trap cleanup EXIT INT TERM HUP

printf 'quote " slash \\\nline' > "$prompt"
"$server" --port "$port" &
server_pid=$!
ready=0
for _attempt in 1 2 3 4 5 6 7 8 9 10; do
    if curl --fail --silent "http://127.0.0.1:$port/health" >/dev/null; then
        ready=1
        break
    fi
    sleep 0.1
done
[ "$ready" -eq 1 ]

GEM_STREAM_MAX_TOKENS=8 "$probe" "http://127.0.0.1:$port/v1" gemma-4-31b \
    "$prompt" "$response" "$summary" >/dev/null

jq -e '
    .schema_version == 1 and .http_code == 200 and
    .requested_max_tokens == 8 and
    .first_visible_event_seconds >= 0.15 and
    .visible_event_count == 2 and .finish_reason == "stop" and
    .usage.total_tokens == 5
' "$summary" >/dev/null

jq -n --rawfile prompt "$prompt" \
    '{model: "gemma-4-31b", messages: [
        {role: "user", content: $prompt},
        {role: "assistant", content: "answer"},
        {role: "user", content: "continue"}], max_tokens: 7}' > "$request"
"$probe" --request "http://127.0.0.1:$port/v1" "$request" \
    "$request_response" "$request_summary" >/dev/null
jq -e '
    .schema_version == 1 and .http_code == 200 and
    .requested_max_tokens == 7 and .visible_event_count == 2 and
    .finish_reason == "stop" and .usage.prompt_tokens == 3
' "$request_summary" >/dev/null