#!/bin/sh
set -eu

[ "$#" -eq 1 ] || { echo "usage: $0 ANALYZER" >&2; exit 2; }
analyzer=$1
fixture=$(mktemp -d /private/tmp/gem-decode-traffic-test.XXXXXX)
trap 'rm -rf "$fixture"' EXIT INT TERM HUP

jq -n '{schema_version:1, engine:"gem-native-c-metal",
  workload:"synthetic-exact-shape-decode-attention-stage-profile",
  shape:"31b-full", value_kernel:"dim32-partitioned",
  value_dimension_tile:32, query_heads:32, kv_heads:4, head_dimension:512,
  tokens:65536, live_tokens:65536, cache_capacity:65536, window:0,
  kv_cache_bytes:1073741824, score_scratch_bytes:8388608,
  value_partition_scratch_bytes:1048576, value_partitions:16}' \
  > "$fixture/profile.json"
"$analyzer" "$fixture/profile.json" "$fixture/analysis.json"
jq -e '
  .schema_version == 1 and
  .workload == "decode-attention-minimum-logical-traffic" and
  .shape.value_dimension_groups == 16 and .shape.value_partitions == 16 and
  .bytes.key_cache_read == 536870912 and
  .bytes.value_cache_read == 536870912 and
  .bytes.raw_score_write == 8388608 and
  .bytes.normalize_score_io == 41943040 and
  .bytes.value_score_read == 134217728 and
  .bytes.avoidable_score_intermediate == 184549376 and
  .bytes.value_partial_write == 1048576 and
  .ratios.score_intermediate_over_kv_cache_read_pct == 17.1875 and
  .interpretation.promotion_eligible == false
' "$fixture/analysis.json" >/dev/null

jq '.score_scratch_bytes += 4' "$fixture/profile.json" > "$fixture/bad-score.json"
if "$analyzer" "$fixture/bad-score.json" >/dev/null 2>&1; then
    echo "traffic analyzer accepted inconsistent score scratch" >&2
    exit 1
fi
jq '.value_partition_scratch_bytes /= 2' "$fixture/profile.json" \
  > "$fixture/bad-partial.json"
if "$analyzer" "$fixture/bad-partial.json" >/dev/null 2>&1; then
    echo "traffic analyzer accepted inconsistent value partial scratch" >&2
    exit 1
fi
jq '.value_partitions = 1.5 | del(.value_partition_scratch_bytes)' \
  "$fixture/profile.json" > "$fixture/fractional.json"
if "$analyzer" "$fixture/fractional.json" >/dev/null 2>&1; then
    echo "traffic analyzer accepted a fractional partition count" >&2
    exit 1
fi
if "$analyzer" "$fixture/profile.json" "$fixture/analysis.json" >/dev/null 2>&1; then
    echo "traffic analyzer overwrote an existing analysis" >&2
    exit 1
fi