#!/bin/sh
set -eu

[ "$#" -eq 1 ] || [ "$#" -eq 2 ] || {
    echo "usage: $0 GATE_OUTPUT_DIR [FROZEN_GEM_BINARY]" >&2
    exit 2
}
directory=$1
binary=${2:-}

for file in gate-config.json print-config.json summary.json prefix-summary.json; do
    [ -s "$directory/$file" ] || {
        echo "gem: missing retained gate artifact: $file" >&2
        exit 1
    }
done

prompt_tokens=$(jq -er '.prompt_tokens' "$directory/gate-config.json")
max_tokens=$(jq -er '.requested_max_tokens' "$directory/gate-config.json")
extension=$(jq -er '.prefix_extension_text | select(length > 0)' \
    "$directory/gate-config.json")

jq -e --slurpfile config "$directory/print-config.json" '
  ($config[0].context | if . < 1024 then . else 1024 end) as $swa_capacity |
  (.schema_version == 1 or .schema_version == 2 or .schema_version == 3 or
   .schema_version == 4) and
  (.schema_version == 1 or
    (.host_policy.preflight_min_tokens >= 0 and
     .host_policy.max_load_per_logical_cpu > 0 and
     (.schema_version < 3 or .host_policy.max_short_window_busy_cpu_percent > 0) and
     (.schema_version < 4 or
       (.host_policy.requires_ac_power == true and
        .host_policy.requires_no_thermal_warning == true and
        .host_policy.monitor_power_and_thermal == true)) and
     .host_policy.monitor_interval_seconds > 0 and
     .host_policy.consecutive_failures_before_abort > 0 and
     .host_policy.sleep_assertion == "caffeinate -i -w server_pid")) and
  .engine == "gem-native-c-metal" and
  .repetitions == 3 and .context == $config[0].context and
  .kv_type == $config[0].kv_type and .model == $config[0].model and
  (.model_sha256 | test("^[0-9a-f]{64}$")) and
  $config[0].engine == "gem-native-c-metal" and
  $config[0].model_size == 31 and $config[0].layers == 60 and
  $config[0].matrix_type == "Q4_0" and $config[0].kv_type == "f32" and
  $config[0].kv_bytes ==
    (50 * $swa_capacity * 16 * 256 * 2 * 4 +
     10 * $config[0].context * 4 * 512 * 2 * 4) and
  $config[0].prefill_batch == 32 and
  $config[0].attention_prefill_rows == 32 and
  ($config[0].prefill_kernel | contains("mma64x8-4sg-at8")) and
  ($config[0].attention_kernel |
    contains("full-decode-value-dim32-split16@8192"))
' "$directory/gate-config.json" >/dev/null

model=$(jq -er '.model' "$directory/gate-config.json")
model_digest_verified=false
if [ -r "$model" ]; then
    expected_model_sha=$(jq -er '.model_sha256' "$directory/gate-config.json")
    actual_model_sha=$(shasum -a 256 "$model" | awk '{print $1}')
    [ "$actual_model_sha" = "$expected_model_sha" ] || {
        echo "gem: retained model digest does not match" >&2
        exit 1
    }
    model_digest_verified=true
fi

prompt_file=$(jq -er '.prompt_file' "$directory/gate-config.json")
prompt_digest_verified=false
if [ -r "$prompt_file" ]; then
    expected_prompt_sha=$(jq -er '.prompt_sha256' "$directory/gate-config.json")
    actual_prompt_sha=$(shasum -a 256 "$prompt_file" | awk '{print $1}')
    [ "$actual_prompt_sha" = "$expected_prompt_sha" ] || {
        echo "gem: retained prompt digest does not match" >&2
        exit 1
    }
    prompt_digest_verified=true
fi

binary_digest_verified=false
binary_sha=
if [ -n "$binary" ]; then
    [ -x "$binary" ] && [ -s "$directory/live-binary-sha256.txt" ] || {
        echo "gem: frozen binary or its digest receipt is missing" >&2
        exit 1
    }
    binary_sha=$(awk 'NR == 1 {print $1}' "$directory/live-binary-sha256.txt")
    actual_binary_sha=$(shasum -a 256 "$binary" | awk '{print $1}')
    [ "$actual_binary_sha" = "$binary_sha" ] || {
        echo "gem: frozen production binary digest changed" >&2
        exit 1
    }
    binary_digest_verified=true
fi

sample=1
gate_schema=$(jq -er '.schema_version' "$directory/gate-config.json")
preflight_min_tokens=$(jq -er '.host_policy.preflight_min_tokens // 0' \
    "$directory/gate-config.json")
while [ "$sample" -le 3 ]; do
    measurement="$directory/measurement-$sample.json"
    prefix_measurement="$directory/prefix-measurement-$sample.json"
    request="$directory/prefix-request-$sample.json"
    log="$directory/server-$sample.log"
    stream="$directory/stream-$sample.sse"
    prefix_stream="$directory/prefix-stream-$sample.sse"
    for file in "$measurement" "$prefix_measurement" "$request" "$log" "$stream" \
                "$prefix_stream"; do
        [ -s "$file" ] || { echo "gem: incomplete retained sample $sample" >&2; exit 1; }
    done
    if [ "$gate_schema" -ge 2 ] && [ "$prompt_tokens" -ge "$preflight_min_tokens" ]; then
        preflight="$directory/host-preflight-$sample.json"
        [ -s "$preflight" ] || {
            echo "gem: missing host preflight receipt for sample $sample" >&2
            exit 1
        }
        jq -e --argjson gate_schema "$gate_schema" '
               .schema_version == (if $gate_schema >= 4 then 2 else 1 end) and
               .status == "pass" and
               .load_per_logical_cpu <= .max_load_per_logical_cpu and
               ($gate_schema < 3 or
                (.short_window_busy_cpu_percent <=
                 .max_short_window_busy_cpu_percent)) and
               ($gate_schema < 4 or
                (.ac_power == true and .no_thermal_warning == true and
                 (.power_state | contains("AC Power"))))' \
            "$preflight" >/dev/null
        [ ! -e "$directory/host-monitor-$sample.failure.txt" ] || {
            echo "gem: retained sample $sample has a host-load monitor failure" >&2
            exit 1
        }
        if [ "$gate_schema" -ge 4 ]; then
            monitor="$directory/host-monitor-$sample.ndjson"
            [ -s "$monitor" ] || {
                echo "gem: missing schema-4 host monitor receipt for sample $sample" >&2
                exit 1
            }
            jq -e -s 'length > 0 and all(
                .schema_version == 2 and .overloaded == false and
                .load_per_logical_cpu <= .max_load_per_logical_cpu and
                .ac_power == true and .no_thermal_warning == true and
                (.power_state | contains("AC Power")))' "$monitor" >/dev/null
        fi
    fi
    jq -e --argjson prompt "$prompt_tokens" --argjson limit "$max_tokens" '
      .schema_version == 1 and .engine == "gem-native-c-metal" and
      .model == "gemma4" and .requested_max_tokens == $limit and
      .http_code == 200 and .first_visible_event_seconds > .header_ttfb_seconds and
      .completion_seconds >= .first_visible_event_seconds and
      .visible_event_count > 0 and .finish_reason != null and
      .usage.prompt_tokens == $prompt and .usage.completion_tokens > 0 and
      .usage.total_tokens == (.usage.prompt_tokens + .usage.completion_tokens)
    ' "$measurement" >/dev/null
    jq -e --argjson prompt "$prompt_tokens" --arg extension "$extension" '
      .model == "gemma4" and .messages[-1].role == "user" and
      .messages[-1].content == $extension
    ' "$request" >/dev/null
    jq -e --argjson prompt "$prompt_tokens" --argjson limit "$max_tokens" '
      .schema_version == 1 and .engine == "gem-native-c-metal" and
      .model == "gemma4" and .requested_max_tokens == $limit and
      .http_code == 200 and .first_visible_event_seconds > .header_ttfb_seconds and
      .completion_seconds >= .first_visible_event_seconds and
      .visible_event_count > 0 and .finish_reason != null and
      .reused_prompt_tokens == $prompt and .appended_prompt_tokens > 0 and
      .usage.prompt_tokens == (.reused_prompt_tokens + .appended_prompt_tokens) and
      .usage.total_tokens == (.usage.prompt_tokens + .usage.completion_tokens)
    ' "$prefix_measurement" >/dev/null
    sed '/^$/d' "$stream" | tail -n 1 | grep -q '^data: \[DONE\]$'
    sed '/^$/d' "$prefix_stream" | tail -n 1 | grep -q '^data: \[DONE\]$'
    stream_usage=$(sed -n '/^data: {/s/^data: //p' "$stream" |
        jq -cS -s '[.[] | select(.usage != null)][-1].usage')
    measured_usage=$(jq -cS '.usage' "$measurement")
    [ "$stream_usage" = "$measured_usage" ] || {
        echo "gem: cold SSE usage differs from measurement $sample" >&2
        exit 1
    }
    stream_finish=$(sed -n '/^data: {/s/^data: //p' "$stream" |
        jq -r -s '[.[] | .choices[]? | select(.finish_reason != null)][-1].finish_reason')
    measured_finish=$(jq -r '.finish_reason' "$measurement")
    [ "$stream_finish" = "$measured_finish" ] || {
        echo "gem: cold SSE finish reason differs from measurement $sample" >&2
        exit 1
    }
    stream_usage=$(sed -n '/^data: {/s/^data: //p' "$prefix_stream" |
        jq -cS -s '[.[] | select(.usage != null)][-1].usage')
    measured_usage=$(jq -cS '.usage' "$prefix_measurement")
    [ "$stream_usage" = "$measured_usage" ] || {
        echo "gem: prefix SSE usage differs from measurement $sample" >&2
        exit 1
    }
    stream_finish=$(sed -n '/^data: {/s/^data: //p' "$prefix_stream" |
        jq -r -s '[.[] | .choices[]? | select(.finish_reason != null)][-1].finish_reason')
    measured_finish=$(jq -r '.finish_reason' "$prefix_measurement")
    [ "$stream_finish" = "$measured_finish" ] || {
        echo "gem: prefix SSE finish reason differs from measurement $sample" >&2
        exit 1
    }
    grep -Fq "gem: prompt prefill $prompt_tokens/$prompt_tokens tokens (100.00%)" "$log"
    appended=$(jq -er '.appended_prompt_tokens' "$prefix_measurement")
    extended=$((prompt_tokens + appended))
    grep -Fq "gem: exact prefix cache hit ($prompt_tokens reused, $appended appended)" "$log"
    grep -Fq "gem: prompt prefill $extended/$extended tokens (100.00%)" "$log"
    raw=$(jq -cS . "$measurement")
    retained=$(jq -cS ".samples[$((sample - 1))]" "$directory/summary.json")
    [ "$raw" = "$retained" ] || {
        echo "gem: cold summary sample $sample differs from raw measurement" >&2
        exit 1
    }
    raw=$(jq -cS . "$prefix_measurement")
    retained=$(jq -cS ".samples[$((sample - 1))]" "$directory/prefix-summary.json")
    [ "$raw" = "$retained" ] || {
        echo "gem: prefix summary sample $sample differs from raw measurement" >&2
        exit 1
    }
    sample=$((sample + 1))
done

canonical_prefix=$(jq -cS . "$directory/prefix-summary.json")
embedded_prefix=$(jq -cS '.prefix_extension' "$directory/summary.json")
[ "$canonical_prefix" = "$embedded_prefix" ] || {
    echo "gem: embedded prefix summary differs from retained prefix summary" >&2
    exit 1
}

jq -n --argjson model_digest_verified "$model_digest_verified" \
    --argjson prompt_digest_verified "$prompt_digest_verified" \
    --argjson binary_digest_verified "$binary_digest_verified" --arg binary_sha "$binary_sha" \
    --slurpfile gate "$directory/gate-config.json" \
    --slurpfile config "$directory/print-config.json" \
    --slurpfile summary "$directory/summary.json" '
  def mean($xs): $xs | add / length;
  def sample_stddev($xs): mean($xs) as $m |
    (($xs | map((. - $m) * (. - $m)) | add) / (($xs | length) - 1)) | sqrt;
  def cv_pct($xs): 100 * sample_stddev($xs) / mean($xs);
  $summary[0] as $s |
  ($s.samples | map(.first_visible_event_seconds)) as $cold_visible |
  ($s.samples | map(.completion_seconds)) as $cold_complete |
  ($s.prefix_extension.samples | map(.first_visible_event_seconds)) as $prefix_visible |
  ($s.prefix_extension.samples | map(.completion_seconds)) as $prefix_complete |
  if (($s.mean_first_visible_event_seconds - mean($cold_visible)) | fabs) > 1e-6 or
      (($s.mean_completion_seconds - mean($cold_complete)) | fabs) > 1e-6 or
      (($s.prefix_extension.mean_first_visible_event_seconds - mean($prefix_visible)) | fabs) > 1e-6 or
      (($s.prefix_extension.mean_completion_seconds - mean($prefix_complete)) | fabs) > 1e-6 or
      $s.repetitions != 3 or $s.prefix_extension.repetitions != 3 or
      $s.prompt_tokens != $gate[0].prompt_tokens or $s.context != $gate[0].context or
      $s.kv_type != $gate[0].kv_type or
      $s.prefix_extension.reused_prompt_tokens != [$gate[0].prompt_tokens] or
      ($s.prefix_extension.appended_prompt_tokens | length) != 1 or
      $s.prefix_extension.appended_prompt_tokens[0] <= 0
  then error("retained summary aggregates are inconsistent") else {
    schema_version: 1,
    result: "pass",
    engine: $gate[0].engine,
    context: $gate[0].context,
    kv_type: $gate[0].kv_type,
    prompt_tokens: $gate[0].prompt_tokens,
    repetitions: 3,
    attention_kernel: $config[0].attention_kernel,
    model_digest_verified: $model_digest_verified,
    prompt_digest_verified: $prompt_digest_verified,
    binary_digest_verified: $binary_digest_verified,
    frozen_binary_sha256: (if $binary_sha == "" then null else $binary_sha end),
    cold: {
      mean_first_visible_event_seconds: mean($cold_visible),
      first_visible_cv_pct: cv_pct($cold_visible),
      mean_completion_seconds: mean($cold_complete),
      completion_cv_pct: cv_pct($cold_complete)
    },
    prefix_extension: {
      mean_first_visible_event_seconds: mean($prefix_visible),
      first_visible_cv_pct: cv_pct($prefix_visible),
      mean_completion_seconds: mean($prefix_complete),
      completion_cv_pct: cv_pct($prefix_complete),
      reused_prompt_tokens: $gate[0].prompt_tokens,
      appended_prompt_tokens: $s.prefix_extension.appended_prompt_tokens[0],
      cold_to_prefix_visible_speedup:
        (mean($cold_visible) / mean($prefix_visible))
    }
  } end
'