#!/bin/sh
set -eu

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

write_full()
{
    path=$1
    position=$2
    model_seconds=$3
    kv_seconds=$4
    jq -n --argjson position "$position" --argjson model_seconds "$model_seconds" \
      --argjson kv_seconds "$kv_seconds" '{
        schema_version:1, engine:"gem-native-c-metal",
        workload:"authentic-one-token-decode-stage-profile", model:"model.gguf",
        attention_kernel:"split16", prefix_mode:"evaluated-prompt",
        attention_selector_override:19, context:($position + 256), position:$position,
        rows:1, repetitions:3, layers:60, mean_gpu_seconds:$model_seconds,
        mean_sampled_gpu_seconds:($model_seconds * 0.99), sample_coverage:0.99,
        gpu_seconds_samples:[$model_seconds,$model_seconds,$model_seconds],
        layers_by_stage:{kv_attention:{mean_seconds:$kv_seconds}}
      }' > "$path"
}

write_operator()
{
    path=$1
    tokens=$2
    score=$3
    normalize=$4
    value=$5
    jq -n --argjson tokens "$tokens" --argjson score "$score" \
      --argjson normalize "$normalize" --argjson value "$value" '
      ($score + $normalize + $value) as $total | {
        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,
        value_partitions:16, query_heads:32, kv_heads:4, head_dimension:512,
        tokens:$tokens, live_tokens:$tokens, cache_capacity:$tokens, window:0,
        repetitions:3, mean_gpu_seconds:$total, mean_sampled_gpu_seconds:($total * 0.999),
        sample_coverage:0.999, gpu_seconds_samples:[$total,$total,$total],
        stages:{score:{mean_seconds:$score}, normalize:{mean_seconds:$normalize},
          value:{mean_seconds:$value}}
      }' > "$path"
}

write_full "$fixture/full64.json" 65536 0.18 0.11
write_operator "$fixture/operator64.json" 65536 0.003 0.0001 0.0029
write_full "$fixture/full128.json" 131072 0.30 0.20
write_operator "$fixture/operator128.json" 131072 0.006 0.0002 0.0058

"$analyzer" "$fixture/full64.json" "$fixture/operator64.json" \
  "$fixture/full128.json" "$fixture/operator128.json" "$fixture/analysis.json"
jq -e '
  .schema_version == 1 and
  .workload == "fused-split-kv-theoretical-ceiling-screen" and
  .fixed_controls.contexts == [65536,131072] and
  .contexts.at_65536.theoretical_whole_token_global_attention_ceiling_pct ==
    33.333333333333336 and
  .contexts.at_65536.theoretical_whole_token_score_only_ceiling_pct ==
    16.666666666666668 and
  .contexts.at_65536.global_attention_gain_required_for_3pct_whole_token_pct == 9 and
  .contexts.at_131072.theoretical_whole_token_global_attention_ceiling_pct == 40 and
  .contexts.at_131072.global_attention_gain_required_for_3pct_whole_token_pct == 7.5 and
  .screening.admissible == true and .screening.promotion_eligible == false and
  (.inputs.at_65536.full_model.sha256 | test("^[0-9a-f]{64}$"))
' "$fixture/analysis.json" >/dev/null

jq '.attention_selector_override = 20' "$fixture/full64.json" > "$fixture/bad-selector.json"
if "$analyzer" "$fixture/bad-selector.json" "$fixture/operator64.json" \
   "$fixture/full128.json" "$fixture/operator128.json" >/dev/null 2>&1; then
    echo "fused ceiling analyzer accepted a non-production selector" >&2
    exit 1
fi
jq '.head_dimension = 256' "$fixture/operator128.json" > "$fixture/bad-shape.json"
if "$analyzer" "$fixture/full64.json" "$fixture/operator64.json" \
   "$fixture/full128.json" "$fixture/bad-shape.json" >/dev/null 2>&1; then
    echo "fused ceiling analyzer accepted a changed operator shape" >&2
    exit 1
fi
jq '.stages.value.mean_seconds *= 0.5' "$fixture/operator64.json" \
  > "$fixture/bad-total.json"
if "$analyzer" "$fixture/full64.json" "$fixture/bad-total.json" \
   "$fixture/full128.json" "$fixture/operator128.json" >/dev/null 2>&1; then
    echo "fused ceiling analyzer accepted inconsistent stage timing" >&2
    exit 1
fi
if "$analyzer" "$fixture/full64.json" "$fixture/operator64.json" \
   "$fixture/full128.json" "$fixture/operator128.json" \
   "$fixture/analysis.json" >/dev/null 2>&1; then
    echo "fused ceiling analyzer overwrote an existing analysis" >&2
    exit 1
fi