#!/bin/sh
set -eu

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

jq -n '{
  schema_version:1, engine:"gem-native-c-metal",
  workload:"d512-fused-split-kv-resource-receipt",
  device:{name:"Apple M4 Max",apple_gpu_family9:true,max_threadgroup_memory_bytes:32768},
  geometry:{query_heads:32,kv_heads:4,heads_per_kv:8,head_dimension:512,
    value_dimension:512,partitions:32,tile_tokens:4,threads_per_threadgroup:256,
    simdgroups_per_threadgroup:8,kv_type:"f32",window:0},
  scratch:{maximum_bytes:4096,denominator_bytes:4096,numerator_bytes:2097152,
    total_bytes:2105344},
  first_pass_pipeline_resources:{thread_execution_width:32,
    max_total_threads_per_threadgroup:1024,static_threadgroup_memory_bytes:0,
    dynamic_threadgroup_memory_bytes:16384,threads_per_threadgroup:256},
  reducer_pipeline_resources:{thread_execution_width:32,
    max_total_threads_per_threadgroup:1024,static_threadgroup_memory_bytes:0,
    dynamic_threadgroup_memory_bytes:0,threads_per_threadgroup:256},
  metal_source_sha256:("a" * 64),native_binary_sha256:("b" * 64)
}' > "$fixture/receipt.json"

"$analyzer" "$fixture/receipt.json" "$fixture/analysis.json"
jq -e '
  .workload == "d512-fused-split-kv-resource-analysis" and
  .compiled_resources.first_pass.total_threadgroup_memory_bytes == 16384 and
  .compiled_resources.first_pass.threadgroup_memory_limit_utilization_pct == 50 and
  .scratch.total_bytes == 2105344 and .screening.passed == true and
  .screening.promotion_eligible == false and
  (.input.sha256 | test("^[0-9a-f]{64}$"))
' "$fixture/analysis.json" >/dev/null

jq '.geometry.tile_tokens = 8 | .first_pass_pipeline_resources.dynamic_threadgroup_memory_bytes = 32768' \
  "$fixture/receipt.json" > "$fixture/bad-tile.json"
if "$analyzer" "$fixture/bad-tile.json" >/dev/null 2>&1; then
    echo "fused split-KV resource analyzer accepted tile8" >&2
    exit 1
fi

jq '.scratch.total_bytes -= 4096' "$fixture/receipt.json" > "$fixture/bad-scratch.json"
if "$analyzer" "$fixture/bad-scratch.json" >/dev/null 2>&1; then
    echo "fused split-KV resource analyzer accepted inconsistent scratch" >&2
    exit 1
fi

jq '.first_pass_pipeline_resources.max_total_threads_per_threadgroup = 128' \
  "$fixture/receipt.json" > "$fixture/bad-threads.json"
if "$analyzer" "$fixture/bad-threads.json" >/dev/null 2>&1; then
    echo "fused split-KV resource analyzer accepted an unlaunchable pipeline" >&2
    exit 1
fi

if "$analyzer" "$fixture/receipt.json" "$fixture/analysis.json" >/dev/null 2>&1; then
    echo "fused split-KV resource analyzer overwrote an existing analysis" >&2
    exit 1
fi