#!/bin/sh
set -eu

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

jq -n '{workload:"d512-online-softmax-cpu-gate-audit",
  screening:{passed:true,promotion_eligible:false}}' > "$fixture/cpu.json"
jq -n '{workload:"d512-fused-split-kv-resource-analysis",
  screening:{passed:true,promotion_eligible:false}}' > "$fixture/resources.json"

write_operator()
{
    tokens=$1
    path=$2
    jq -n --argjson tokens "$tokens" '{schema_version:1,engine:"gem-native-c-metal",
      workload:"d512-fused-split-kv-operator-differential",tokens:$tokens,
      query_heads:32,kv_heads:4,head_dimension:512,value_dimension:512,partitions:32,
      tile_tokens:4,kv_type:"f32",window:0,baseline_selector:19,
      candidate_route:"d512-fused-split-kv-p32-tile4",repetitions:3,finite:true,
      deterministic:true,correlation:0.999999,nmse:0.000001,maximum_absolute_error:0.00001,
      candidate_hash:"0123456789abcdef",candidate_repeat_hash:"0123456789abcdef",passed:true
    }' > "$path"
}

write_model()
{
    position=$1
    path=$2
    jq -n --argjson position "$position" '{schema_version:1,engine:"gem-native-c-metal",
      workload:"d512-fused-split-kv-model-differential",model_size:31,layers:60,
      model_sha256:("a"*64),prefix_mode:"evaluated-prompt",position:$position,
      context:($position+256),baseline_selector:19,
      candidate_route:"d512-fused-split-kv-p32-tile4",production_unchanged:true,
      finite:true,deterministic:true,correlation:0.999999,nmse:0.000001,
      baseline_argmax:42,candidate_argmax:42,
      baseline_top10:[42,1,2,3,4,5,6,7,8,9],candidate_top10:[42,1,2,3,4,5,6,7,8,9],
      baseline_continuation_tokens:[11,12,13,14,15,16,17,18],
      candidate_continuation_tokens:[11,12,13,14,15,16,17,18],
      candidate_logits_hash:"fedcba9876543210",
      candidate_repeat_logits_hash:"fedcba9876543210",passed:true
    }' > "$path"
}

for tokens in 8191 8192 8193 65536 131072; do
    write_operator "$tokens" "$fixture/op-$tokens.json"
done
for position in 8191 8192 8193; do
    write_model "$position" "$fixture/model-$position.json"
done

"$analyzer" "$fixture/cpu.json" "$fixture/resources.json" \
  "$fixture/op-8191.json" "$fixture/op-8192.json" "$fixture/op-8193.json" \
  "$fixture/op-65536.json" "$fixture/op-131072.json" \
  "$fixture/model-8191.json" "$fixture/model-8192.json" "$fixture/model-8193.json" \
  "$fixture/analysis.json"
jq -e '.workload=="d512-fused-split-kv-correctness-analysis" and
  .fixed_controls.operator_tokens==[8191,8192,8193,65536,131072] and
  .screening.passed==true and .screening.promotion_eligible==false and
  (.inputs.models.at_8192.sha256|test("^[0-9a-f]{64}$"))' \
  "$fixture/analysis.json" >/dev/null

jq '.candidate_continuation_tokens[-1]=99' "$fixture/model-8192.json" \
  > "$fixture/bad-continuation.json"
if "$analyzer" "$fixture/cpu.json" "$fixture/resources.json" \
   "$fixture/op-8191.json" "$fixture/op-8192.json" "$fixture/op-8193.json" \
   "$fixture/op-65536.json" "$fixture/op-131072.json" \
   "$fixture/model-8191.json" "$fixture/bad-continuation.json" \
   "$fixture/model-8193.json" >/dev/null 2>&1; then
    echo "fused split-KV correctness analyzer accepted continuation divergence" >&2
    exit 1
fi

jq '.tokens=8194' "$fixture/op-8193.json" > "$fixture/bad-boundary.json"
if "$analyzer" "$fixture/cpu.json" "$fixture/resources.json" \
   "$fixture/op-8191.json" "$fixture/op-8192.json" "$fixture/bad-boundary.json" \
   "$fixture/op-65536.json" "$fixture/op-131072.json" \
   "$fixture/model-8191.json" "$fixture/model-8192.json" "$fixture/model-8193.json" \
   >/dev/null 2>&1; then
    echo "fused split-KV correctness analyzer accepted changed boundary" >&2
    exit 1
fi

if "$analyzer" "$fixture/cpu.json" "$fixture/resources.json" \
   "$fixture/op-8191.json" "$fixture/op-8192.json" "$fixture/op-8193.json" \
   "$fixture/op-65536.json" "$fixture/op-131072.json" \
   "$fixture/model-8191.json" "$fixture/model-8192.json" "$fixture/model-8193.json" \
   "$fixture/analysis.json" >/dev/null 2>&1; then
    echo "fused split-KV correctness analyzer overwrote existing analysis" >&2
    exit 1
fi