#!/bin/sh
# Isolated correctness, ABBA operator, and long-prefill-attribution gate for selector 20.
set -eu

if [ "$#" -ne 2 ]; then
    echo "usage: $0 MODEL.gguf OUTPUT_DIR" >&2
    exit 2
fi
model=$1
output=$2
build_dir=${GEM_SCORE_TILE8_BUILD:-build-score-tile8}
backend_ops=${GEM_LLAMA_BACKEND_OPS:-../qwen38-llama-bench/llama.cpp/build-gem-metal/bin/test-backend-ops}
max_load_per_cpu=${GEM_SCORE_TILE8_MAX_LOAD_PER_CPU:-1.0}
max_short_busy=${GEM_SCORE_TILE8_MAX_SHORT_WINDOW_BUSY_CPU_PERCENT:-10.0}

[ -r "$model" ] || { echo "gem: model is not readable: $model" >&2; exit 2; }
[ ! -e "$output" ] || { echo "gem: output already exists: $output" >&2; exit 2; }
for tool in jq shasum otool sw_vers sysctl pmset top; do
    command -v "$tool" >/dev/null 2>&1 || {
        echo "gem: required score-tile8 gate tool is missing: $tool" >&2
        exit 2
    }
done
[ "$(uname -s)" = Darwin ] || { echo "gem: score-tile8 gate requires macOS" >&2; exit 2; }
for value in "$max_load_per_cpu" "$max_short_busy"; do
    awk -v value="$value" 'BEGIN {exit !(value ~ /^[0-9]+([.][0-9]+)?$/ && value > 0)}' || {
        echo "gem: score-tile8 host thresholds must be positive numbers" >&2
        exit 2
    }
done
if command -v pgrep >/dev/null 2>&1 && pgrep -f 'build/gem serve' >/dev/null 2>&1; then
    echo "gem: refusing to overlap a production server Metal workload" >&2
    exit 1
fi
model_sha256=${GEM_MODEL_SHA256:-}
if [ -n "$model_sha256" ]; then
    case $model_sha256 in *[!0-9a-fA-F]*) model_sha256=invalid ;; esac
    [ "${#model_sha256}" -eq 64 ] || {
        echo "gem: GEM_MODEL_SHA256 must contain exactly 64 hexadecimal digits" >&2
        exit 2
    }
else
    model_sha256=$(shasum -a 256 "$model" | awk '{print $1}')
fi
for binary in "$build_dir/gem" "$build_dir/gem-attention-split-test" \
              "$build_dir/gem-decode-attention-profile" \
              "$build_dir/gem-long-prefill-stage-profile" "$backend_ops"; do
    [ -x "$binary" ] || { echo "gem: required binary is missing: $binary" >&2; exit 2; }
done

mkdir -p "$output"
printf '%s  %s\n' "$model_sha256" "$model" > "$output/model-sha256.txt"
date -u '+%Y-%m-%dT%H:%M:%SZ' > "$output/run-start-utc.txt"
sw_vers > "$output/macos-version.txt"
shasum -a 256 "$build_dir/gem" "$build_dir/gem-attention-split-test" \
    "$build_dir/gem-decode-attention-profile" \
    "$build_dir/gem-long-prefill-stage-profile" src/metal.m src/metal_layer.c \
    src/decode_attention_profile.c scripts/run-score-tile8-gate.sh \
    scripts/analyze-score-tile8-decision.sh > "$output/binary-sha256.txt"
"$build_dir/gem" print-config --model "$model" --context 131328 --kv-type f32 \
    > "$output/print-config.json"
jq -e '
  .engine == "gem-native-c-metal" and .model_size == 31 and .layers == 60 and
  .context == 131328 and .kv_type == "f32" and
  (.attention_kernel | contains("full-decode-value-dim32-split16@8192"))
' "$output/print-config.json" >/dev/null
otool -L "$build_dir/gem" "$build_dir/gem-attention-split-test" \
    "$build_dir/gem-decode-attention-profile" \
    "$build_dir/gem-long-prefill-stage-profile" > "$output/native-linkage.txt"
if grep -Eiq 'lib(llama|ggml)' "$output/native-linkage.txt"; then
    echo "gem: diagnostic binaries unexpectedly link llama/ggml" >&2
    exit 1
fi
jq -n --argjson max_load "$max_load_per_cpu" --argjson max_busy "$max_short_busy" '
  {schema_version: 2, engine: "gem-native-c-metal",
   axis: {name: "score_token_tile", baseline: 1, candidate: 8},
   contexts: [65536, 131072], repetitions_per_arm: 3, iterations: 5, warmup: 3,
   abba_order: ["baseline-a1", "candidate-b1", "candidate-b2", "baseline-a2"],
   requires_ac_power: true, requires_no_thermal_warning: true,
   max_load_per_logical_cpu: $max_load,
   max_short_window_busy_cpu_percent: $max_busy,
   correctness_after_timing: true}' > "$output/gate-config.json"

host_receipt() {
    phase=$1
    tokens=$2
    destination="$output/host-$phase-$tokens.json"
    timestamp=$(date -u '+%Y-%m-%dT%H:%M:%SZ')
    load_average=$(sysctl -n vm.loadavg)
    load_1m=$(printf '%s\n' "$load_average" | awk '{print $2}')
    logical_cpus=$(sysctl -n hw.logicalcpu)
    load_per_cpu=$(awk -v load="$load_1m" -v cpus="$logical_cpus" \
        'BEGIN {printf "%.6f", load / cpus}')
    recent_cpu=$(ps -Ao pcpu= | awk '{sum += $1} END {printf "%.3f", sum}')
    short_busy=$(LC_ALL=C top -l 2 -s 1 -n 0 -stats pid,cpu,command |
        awk -F'[:,%]' '/^CPU usage:/ {idle=$6} END {
          gsub(/^[[:space:]]+|[[:space:]]+$/, "", idle)
          if (idle !~ /^[0-9]+([.][0-9]+)?$/) exit 1
          printf "%.3f", 100 - idle
        }') || return 1
    power=$(LC_ALL=C pmset -g batt 2>&1 || true)
    thermal=$(LC_ALL=C pmset -g therm 2>&1 || true)
    ac=false; case $power in *"Now drawing from 'AC Power'"*) ac=true ;; esac
    no_warning=false
    if printf '%s\n' "$thermal" | grep -Fq 'No thermal warning level has been recorded' &&
       printf '%s\n' "$thermal" | grep -Fq 'No performance warning level has been recorded'; then
        no_warning=true
    elif printf '%s\n' "$thermal" | grep -Eq 'CPU_Scheduler_Limit[[:space:]]*=[[:space:]]*100' &&
         printf '%s\n' "$thermal" | grep -Eq 'CPU_Speed_Limit[[:space:]]*=[[:space:]]*100'; then
        no_warning=true
    fi
    status=pass
    awk -v actual="$load_per_cpu" -v maximum="$max_load_per_cpu" \
        'BEGIN {exit !(actual <= maximum)}' || status=fail
    awk -v actual="$short_busy" -v maximum="$max_short_busy" \
        'BEGIN {exit !(actual <= maximum)}' || status=fail
    [ "$ac" = true ] || status=fail
    [ "$no_warning" = true ] || status=fail
    top_processes=$(ps -Ao pid,pcpu,pmem,etime,comm -r | sed -n '1,24p')
    jq -n --arg timestamp "$timestamp" --arg phase "$phase" --arg status "$status" \
      --arg load_average "$load_average" --argjson load_1m "$load_1m" \
      --argjson logical_cpus "$logical_cpus" --argjson load_per_cpu "$load_per_cpu" \
      --argjson max_load "$max_load_per_cpu" --argjson recent_cpu "$recent_cpu" \
      --argjson short_busy "$short_busy" --argjson max_busy "$max_short_busy" \
      --arg power "$power" --argjson ac "$ac" --arg thermal "$thermal" \
      --argjson no_warning "$no_warning" --arg top_processes "$top_processes" '
      {schema_version: 2, timestamp_utc: $timestamp, phase: $phase, status: $status,
       load_average: $load_average, load_1m: $load_1m, logical_cpus: $logical_cpus,
       load_per_logical_cpu: $load_per_cpu, max_load_per_logical_cpu: $max_load,
       recent_aggregate_process_cpu_percent: $recent_cpu,
       short_window_busy_cpu_percent: $short_busy,
       max_short_window_busy_cpu_percent: $max_busy,
       power_state: $power, ac_power: $ac, thermal_state: $thermal,
       no_thermal_warning: $no_warning, top_processes: $top_processes}' > "$destination"
    [ "$status" = pass ]
}

run_profile() {
    kernel=$1
    tokens=$2
    destination=$3
    temporary="$destination.tmp"
    "$build_dir/gem-decode-attention-profile" --shape 31b-full --kernel "$kernel" \
        --tokens "$tokens" --iterations 5 --repetitions 3 --warmup 3 > "$temporary"
    jq -e --arg kernel "$kernel" --argjson tokens "$tokens" '
      .schema_version == 1 and
      .workload == "synthetic-exact-shape-decode-attention-stage-profile" and
      .shape == "31b-full" and .tokens == $tokens and .live_tokens == $tokens and
      .repetitions == 3 and .iterations == 5 and .warmup == 3 and
      .value_kernel == "dim32-partitioned" and .value_partitions == 16 and
      (if $kernel == "dim32-split16" then .score_token_tile == 1
       else .score_token_tile == 8 end)
    ' "$temporary" >/dev/null
    mv "$temporary" "$destination"
}

for tokens in 65536 131072; do
    if ! host_receipt preflight "$tokens"; then
        echo "gem: score-tile8 gate rejected busy host before $tokens-token screen" >&2
        exit 1
    fi
    run_profile dim32-split16 "$tokens" "$output/score-a1-$tokens.json"
    run_profile dim32-split16-score-tile8 "$tokens" "$output/score-b1-$tokens.json"
    run_profile dim32-split16-score-tile8 "$tokens" "$output/score-b2-$tokens.json"
    run_profile dim32-split16 "$tokens" "$output/score-a2-$tokens.json"
    scripts/analyze-decode-attention-abba.sh \
        "$output/score-a1-$tokens.json" "$output/score-b1-$tokens.json" \
        "$output/score-b2-$tokens.json" "$output/score-a2-$tokens.json" \
        "$output/score-comparison-$tokens.json"
    if ! host_receipt postflight "$tokens"; then
        echo "gem: score-tile8 gate rejected changed host after $tokens-token screen" >&2
        exit 1
    fi
done

# Correctness work follows the timing arms so it cannot warm or load-bias them.
"$backend_ops" test -b MTL0 -o FLASH_ATTN_EXT \
    > "$output/adjacent-flash-attn-backend-ops-mtl0.txt" 2>&1
oracle_counts=$(sed -n 's/.* \([0-9][0-9]*\)\/\([0-9][0-9]*\) tests passed.*/\1 \2/p' \
    "$output/adjacent-flash-attn-backend-ops-mtl0.txt" | tail -n 1)
oracle_passed=$(printf '%s\n' "$oracle_counts" | awk '{print $1}')
oracle_total=$(printf '%s\n' "$oracle_counts" | awk '{print $2}')
if [ -z "$oracle_passed" ] || [ -z "$oracle_total" ] ||
   [ "$oracle_passed" -ne "$oracle_total" ] || [ "$oracle_passed" -eq 0 ]; then
    echo "gem: adjacent FLASH_ATTN_EXT oracle did not fully pass" >&2
    exit 1
fi
printf '{"passed":%s,"total":%s}\n' "$oracle_passed" "$oracle_total" \
    > "$output/adjacent-oracle-summary.json"
"$build_dir/gem-attention-split-test" > "$output/native-attention-parity.txt" 2>&1

"$build_dir/gem-long-prefill-stage-profile" "$model" 131040 3 --seed-prefix \
    > "$output/prefill-stage-seeded-131040.json"
scripts/analyze-prefill-stages.sh "$output/prefill-stage-seeded-131040.json" \
    "$output/prefill-stage-analysis-131040.json"

scripts/analyze-score-tile8-decision.sh \
    "$output/score-comparison-65536.json" "$output/score-comparison-131072.json" \
    "$output/prefill-stage-analysis-131040.json" "$output/adjacent-oracle-summary.json" \
    "$model" "$model_sha256" \
    "$output/gate-config.json" \
    "$output/host-preflight-65536.json" "$output/host-postflight-65536.json" \
    "$output/host-preflight-131072.json" "$output/host-postflight-131072.json" \
    "$output/decision.json"
jq . "$output/decision.json"