#!/bin/sh
# Authentic native-server disconnect during prompt prefill, followed by same-process recovery.
set -eu

if [ "$#" -ne 5 ]; then
    echo "usage: $0 MODEL.gguf CONTEXT PROMPT_FILE OUTPUT_DIR PORT" >&2
    exit 2
fi

model=$1
context=$2
prompt=$3
output_dir=$4
port=$5
cancel_after=${GEM_CANCEL_AFTER_SECONDS:-5}
recovery_limit=${GEM_CANCEL_RECOVERY_LIMIT_SECONDS:-60}
kv_type=${GEM_KV_TYPE:-f32}
case $context:$port:$cancel_after:$recovery_limit in
    *[!0-9:]*|:*|*::*)
        echo "gem: context, port, and cancellation limits must be positive integers" >&2
        exit 2
        ;;
esac
if [ "$context" -eq 0 ] || [ "$port" -eq 0 ] || [ "$cancel_after" -eq 0 ] ||
   [ "$recovery_limit" -eq 0 ]; then
    echo "gem: context, port, and cancellation limits must be positive integers" >&2
    exit 2
fi
if [ "$port" -gt 65535 ]; then
    echo "gem: port must not exceed 65535" >&2
    exit 2
fi
case $kv_type in
    f32|f16|f16-full) ;;
    *) echo "gem: GEM_KV_TYPE must be f32, f16, or f16-full" >&2; exit 2 ;;
esac
if [ ! -r "$model" ] || [ ! -r "$prompt" ]; then
    echo "gem: model and prompt must be readable" >&2
    exit 2
fi
if [ -e "$output_dir" ]; then
    echo "gem: output directory already exists: $output_dir" >&2
    exit 2
fi
if command -v lsof >/dev/null 2>&1 &&
   lsof -nP -iTCP:"$port" -sTCP:LISTEN >/dev/null 2>&1; then
    echo "gem: cancellation-gate port $port is already occupied" >&2
    exit 1
fi

prompt_count_json=$(build/gem-tokenize --count --file --chat-user "$model" "$prompt")
prompt_tokens=$(printf '%s\n' "$prompt_count_json" | jq -er '.token_count')
if [ "$prompt_tokens" -lt 128 ] || [ "$prompt_tokens" -gt $((context - 2)) ]; then
    echo "gem: cancellation prompt must contain 128..$((context - 2)) chat tokens" >&2
    exit 2
fi
model_sha256=${GEM_MODEL_SHA256:-}
if [ -z "$model_sha256" ]; then
    model_sha256=$(shasum -a 256 "$model" | awk '{print $1}')
fi
case $model_sha256 in
    *[!0-9a-fA-F]*) echo "gem: invalid model SHA-256" >&2; exit 2 ;;
esac
if [ "${#model_sha256}" -ne 64 ]; then
    echo "gem: invalid model SHA-256" >&2
    exit 2
fi
prompt_sha256=$(shasum -a 256 "$prompt" | awk '{print $1}')

mkdir -p "$output_dir"
build/gem print-config --model "$model" --context "$context" --kv-type "$kv_type" \
    > "$output_dir/print-config.json"
jq -e --arg kv_type "$kv_type" --argjson context "$context" \
    '.engine == "gem-native-c-metal" and .context == $context and .kv_type == $kv_type' \
    "$output_dir/print-config.json" >/dev/null
jq -n --arg model "$model" --arg model_sha256 "$model_sha256" \
    --arg prompt "$prompt" --arg prompt_sha256 "$prompt_sha256" --arg kv_type "$kv_type" \
    --argjson context "$context" --argjson prompt_tokens "$prompt_tokens" \
    --argjson cancel_after_seconds "$cancel_after" \
    --argjson recovery_limit_seconds "$recovery_limit" \
    '{schema_version: 1, engine: "gem-native-c-metal", model: $model,
      model_sha256: $model_sha256, prompt_file: $prompt, prompt_sha256: $prompt_sha256,
      context: $context, kv_type: $kv_type, prompt_tokens: $prompt_tokens,
      cancel_after_seconds: $cancel_after_seconds,
      recovery_limit_seconds: $recovery_limit_seconds}' > "$output_dir/gate-config.json"
jq -n --rawfile prompt "$prompt" \
    '{model: "gemma4", stream: true, max_tokens: 2,
      messages: [{role: "user", content: $prompt}]}' > "$output_dir/cancel-request.json"
printf 'Reply with exactly OK.\n' > "$output_dir/recovery-prompt.txt"

server_pid=
cleanup()
{
    if [ -n "$server_pid" ] && kill -0 "$server_pid" 2>/dev/null; then
        kill "$server_pid"
        wait "$server_pid" 2>/dev/null || true
    fi
    server_pid=
}
trap cleanup EXIT INT TERM HUP

build/gem serve --model "$model" --alias gemma4 --port "$port" \
    --context "$context" --max-tokens 2 --command-span token --kv-type "$kv_type" \
    > "$output_dir/server.stdout" 2> "$output_dir/server.stderr" &
server_pid=$!
ready=0
attempt=0
while [ "$attempt" -lt 1200 ]; do
    if curl --fail --silent --max-time 1 \
        "http://127.0.0.1:$port/health" >/dev/null 2>&1; then
        ready=1
        break
    fi
    if ! kill -0 "$server_pid" 2>/dev/null; then
        echo "gem: server exited before cancellation-gate readiness" >&2
        exit 1
    fi
    attempt=$((attempt + 1))
    sleep 0.1
done
if [ "$ready" -ne 1 ]; then
    echo "gem: cancellation-gate readiness timed out" >&2
    exit 1
fi

set +e
curl --silent --show-error --no-buffer --max-time "$cancel_after" \
    -H 'Content-Type: application/json' --data-binary @"$output_dir/cancel-request.json" \
    "http://127.0.0.1:$port/v1/chat/completions" \
    > "$output_dir/cancelled.sse" 2> "$output_dir/curl.stderr"
curl_status=$?
set -e
if [ "$curl_status" -ne 28 ]; then
    echo "gem: cancellation curl returned $curl_status, expected timeout status 28" >&2
    exit 1
fi
keepalives=$(grep -c '^: keepalive$' "$output_dir/cancelled.sse" || true)
if [ "$keepalives" -lt 1 ]; then
    echo "gem: cancellation stream contained no keepalive" >&2
    exit 1
fi
if grep -Fq "gem: prompt prefill $prompt_tokens/$prompt_tokens tokens (100.00%)" \
    "$output_dir/server.stderr"; then
    echo "gem: cancellation prompt completed before the disconnect was observed" >&2
    exit 1
fi

recovery_start=$(date +%s)
health_status=0
while :; do
    if curl --fail --silent --max-time 1 \
        "http://127.0.0.1:$port/health" >/dev/null 2>&1; then
        health_status=200
        break
    fi
    recovery_now=$(date +%s)
    if [ $((recovery_now - recovery_start)) -ge "$recovery_limit" ]; then
        echo "gem: server did not recover within $recovery_limit seconds" >&2
        exit 1
    fi
    if ! kill -0 "$server_pid" 2>/dev/null; then
        echo "gem: server exited after cancellation" >&2
        exit 1
    fi
    sleep 0.1
done
recovery_end=$(date +%s)
recovery_seconds=$((recovery_end - recovery_start))

GEM_STREAM_MAX_TOKENS=16 scripts/server-stream-probe.sh \
    "http://127.0.0.1:$port/v1" gemma4 "$output_dir/recovery-prompt.txt" \
    "$output_dir/recovery.sse" "$output_dir/recovery-measurement.json"
jq -e '.http_code == 200 and .requested_max_tokens == 16 and
       .visible_event_count > 0 and .finish_reason != null and
       .usage.prompt_tokens > 0 and .usage.completion_tokens > 0 and
       .usage.total_tokens == (.usage.prompt_tokens + .usage.completion_tokens)' \
    "$output_dir/recovery-measurement.json" >/dev/null

jq -n --arg model "$model" --arg model_sha256 "$model_sha256" \
    --arg prompt_sha256 "$prompt_sha256" --arg kv_type "$kv_type" \
    --argjson context "$context" --argjson prompt_tokens "$prompt_tokens" \
    --argjson curl_status "$curl_status" --argjson keepalives "$keepalives" \
    --argjson health_status "$health_status" --argjson recovery_seconds "$recovery_seconds" \
    '{schema_version: 1, engine: "gem-native-c-metal", workload: "prefill-cancel-recovery",
      model: $model, model_sha256: $model_sha256, prompt_sha256: $prompt_sha256,
      context: $context, kv_type: $kv_type, prompt_tokens: $prompt_tokens,
      cancellation_curl_status: $curl_status, keepalive_events: $keepalives,
      post_cancel_health_status: $health_status,
      post_cancel_health_recovery_seconds: $recovery_seconds,
      recovery_requested_max_tokens: 16,
      same_process_recovery_request: "passed"}' > "$output_dir/summary.json"
jq . "$output_dir/summary.json"