#!/bin/sh
# Measure the first visible OpenAI SSE delta, not the immediate HTTP headers/role event.
set -eu

if [ "$#" -ne 5 ]; then
    echo "usage: $0 BASE_URL MODEL PROMPT_FILE OUTPUT_SSE OUTPUT_JSON" >&2
    echo "       $0 --request BASE_URL REQUEST_JSON OUTPUT_SSE OUTPUT_JSON" >&2
    exit 1
fi

max_tokens=${GEM_STREAM_MAX_TOKENS:-128}
case $max_tokens in
    ''|*[!0-9]*|0) echo "gem: GEM_STREAM_MAX_TOKENS must be a positive integer" >&2; exit 2 ;;
esac

if [ "$1" = "--request" ]; then
    base_url=$2
    request_file=$3
    output_sse=$4
    output_json=$5
    if [ ! -r "$request_file" ]; then
        echo "gem: request file is not readable: $request_file" >&2
        exit 1
    fi
    request=$(jq -ce '
        if (.model | type) != "string" or (.messages | type) != "array" then
          error("request must contain string model and array messages")
        else
          .stream = true |
          .stream_options = ((.stream_options // {}) + {include_usage: true})
        end
    ' "$request_file")
    model=$(printf '%s\n' "$request" | jq -r '.model')
    max_tokens=$(printf '%s\n' "$request" | jq -r '.max_tokens // 128')
else
    base_url=$1
    model=$2
    prompt_file=$3
    output_sse=$4
    output_json=$5
    if [ ! -r "$prompt_file" ]; then
        echo "gem: prompt file is not readable: $prompt_file" >&2
        exit 1
    fi
    request=$(jq -n --rawfile prompt "$prompt_file" --arg model "$model" \
        --argjson max_tokens "$max_tokens" \
        '{model: $model, messages: [{role: "user", content: $prompt}], max_tokens: $max_tokens,
          stream: true, stream_options: {include_usage: true}}')
fi
trace=$(mktemp /private/tmp/gem-server-stream-trace.XXXXXX)
metrics=$(mktemp /private/tmp/gem-server-stream-metrics.XXXXXX)
events=$(mktemp /private/tmp/gem-server-stream-events.XXXXXX)
trap 'rm -f "$trace" "$metrics" "$events"' EXIT

curl --fail --silent --show-error --no-buffer --trace-time --trace-ascii "$trace" \
    --output "$output_sse" --write-out '%{time_starttransfer} %{time_total} %{http_code}\n' \
    -H 'Content-Type: application/json' --data "$request" \
    "$base_url/chat/completions" > "$metrics"

ttft=$(awk '
function seconds(value, parts) {
    split(value, parts, ":")
    return parts[1] * 3600 + parts[2] * 60 + parts[3]
}
function inspect() {
    if (receiving && visible == "" &&
        (payload ~ /"content":[[:space:]]*"[^"]/ ||
         payload ~ /"reasoning_content":[[:space:]]*"[^"]/ ||
         payload ~ /"tool_calls":[[:space:]]*\[[[:space:]]*\{/))
        visible = received_at
}
{
    if (started == "") started = seconds($1)
    if ($0 ~ /<= Recv data,/) {
        inspect()
        receiving = 1
        received_at = seconds($1)
        payload = ""
        next
    }
    if ($1 ~ /^[0-9][0-9]:[0-9][0-9]:/) {
        inspect()
        receiving = 0
    }
    if (receiving && $0 ~ /^[0-9a-fA-F][0-9a-fA-F][0-9a-fA-F][0-9a-fA-F]: /)
        payload = payload substr($0, 7)
}
END {
    inspect()
    if (visible == "") exit 1
    delta = visible - started
    if (delta < 0) delta += 86400
    printf "%.6f", delta
}
' "$trace") || {
    echo "gem: stream contained no visible content, reasoning, or tool-call delta" >&2
    exit 1
}

sed -n 's/^data: //p' "$output_sse" | sed '/^\[DONE\]$/d' > "$events"
event_summary=$(jq -sc '
    def nonempty_string: type == "string" and length > 0;
    def nonempty_tool_calls: type == "array" and length > 0;
    {
      visible_event_count: map(select(
        (.choices[0].delta.content? | nonempty_string) or
        (.choices[0].delta.reasoning_content? | nonempty_string) or
        (.choices[0].delta.tool_calls? | nonempty_tool_calls))) | length,
      finish_reason: ([.[] | .choices[0].finish_reason? | select(. != null)] | last),
      usage: ([.[] | .usage? | select(. != null)] | last)
    }
' "$events")
header_ttfb=$(awk '{print $1}' "$metrics")
total=$(awk '{print $2}' "$metrics")
http_code=$(awk '{print $3}' "$metrics")

jq -n \
    --arg engine "gem-native-c-metal" \
    --arg model "$model" \
    --argjson requested_max_tokens "$max_tokens" \
    --argjson header_ttfb "$header_ttfb" \
    --argjson visible_ttft "$ttft" \
    --argjson total "$total" \
    --argjson http_code "$http_code" \
    --argjson events "$event_summary" \
    '{schema_version: 1, engine: $engine, model: $model,
      requested_max_tokens: $requested_max_tokens,
      header_ttfb_seconds: $header_ttfb,
      first_visible_event_seconds: $visible_ttft,
      completion_seconds: $total,
      http_code: $http_code,
      visible_event_count: $events.visible_event_count,
      finish_reason: $events.finish_reason,
      usage: $events.usage}' > "$output_json"

jq -e '
    .http_code == 200 and .first_visible_event_seconds > .header_ttfb_seconds and
    # curl and trace-ascii use independent clocks/precision; permit sub-millisecond ordering skew.
    .completion_seconds + 0.001 >= .first_visible_event_seconds and
    .visible_event_count > 0 and .finish_reason != null and .usage.total_tokens > 0
' "$output_json" >/dev/null
jq . "$output_json"