#!/bin/sh
# Quiet-host, one-axis BQ8/BK64 versus BQ8/BK128 operator gate.
set -eu

[ "$#" -eq 1 ] || {
    echo "usage: $0 OUTPUT_DIRECTORY" >&2
    exit 2
}
output=$1
build_dir=${GEM_BQ8_GATE_BUILD:-build}
profile="$build_dir/gem-long-attention-profile"
counter_info="$build_dir/gem-metal-counter-info"
max_load_per_cpu=${GEM_BQ8_GATE_MAX_LOAD_PER_CPU:-1.0}
max_short_window_busy_cpu=${GEM_BQ8_GATE_MAX_SHORT_WINDOW_BUSY_CPU_PERCENT:-10.0}
repetitions=10
warmup=10

[ ! -e "$output" ] || { echo "gem: output already exists: $output" >&2; exit 2; }
[ -x "$profile" ] || { echo "gem: required profiler is missing: $profile" >&2; exit 2; }
[ -x "$counter_info" ] || {
    echo "gem: required Metal inventory tool is missing: $counter_info" >&2
    exit 2
}
case $max_load_per_cpu in
    ''|*[!0-9.]*) echo "gem: invalid maximum load per CPU" >&2; exit 2 ;;
esac
awk -v maximum="$max_load_per_cpu" 'BEGIN {exit !(maximum > 0)}' || {
    echo "gem: maximum load per CPU must be positive" >&2
    exit 2
}
awk -v maximum="$max_short_window_busy_cpu" 'BEGIN {exit !(maximum > 0)}' || {
    echo "gem: maximum short-window busy CPU must be positive" >&2
    exit 2
}
if [ "$(uname -s)" != Darwin ] || ! command -v sysctl >/dev/null 2>&1 ||
   ! command -v pmset >/dev/null 2>&1 || ! command -v caffeinate >/dev/null 2>&1 ||
   ! command -v top >/dev/null 2>&1; then
    echo "gem: BK128 performance gate requires macOS host-state tools" >&2
    exit 2
fi
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

mkdir -p "$output"
date -u '+%Y-%m-%dT%H:%M:%SZ' > "$output/run-start-utc.txt"
sw_vers > "$output/macos-version.txt"
shasum -a 256 "$profile" "$counter_info" src/metal.m src/long_attention_profile.c \
    scripts/run-bq8-bk128-gate.sh scripts/analyze-bq8-bk128-gate.sh \
    > "$output/sha256.txt"
otool -L "$profile" "$counter_info" > "$output/native-linkage.txt"
if grep -Eiq 'lib(llama|ggml)' "$output/native-linkage.txt"; then
    echo "gem: profiler unexpectedly links llama/ggml" >&2
    exit 1
fi
"$counter_info" > "$output/metal-device.json"
jq -e '
  .schema_version == 1 and .workload == "metal-counter-info" and
  .device.name == "Apple M4 Max" and .device.has_unified_memory == true and
  .device.gpu_families.apple9 == true and
  .device.max_threadgroup_memory_bytes >= 20544
' "$output/metal-device.json" >/dev/null
jq -n --arg build_dir "$build_dir" --argjson maximum "$max_load_per_cpu" \
    --argjson max_short_window_busy_cpu "$max_short_window_busy_cpu" \
    --argjson repetitions "$repetitions" --argjson warmup "$warmup" \
    '{schema_version: 2, engine: "gem-native-c-metal",
      axis: {name: "cache_tile_width", baseline: 64, candidate: 128},
      positions: [4064, 65504, 131040], rows: 32, kv_type: "f32",
      repetitions: $repetitions, warmup: $warmup, seed_prefix: true,
      max_load_per_logical_cpu: $maximum,
      max_short_window_busy_cpu_percent: $max_short_window_busy_cpu,
      requires_ac_power: true, requires_no_thermal_warning: true,
      build_directory: $build_dir}' \
    > "$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_process_cpu_percent=$(ps -Ao pcpu= |
        awk '{sum += $1} END {printf "%.3f", sum}')
    short_window_busy_cpu=$(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
        }') || {
        echo "gem: could not sample short-window macOS CPU state" >&2
        return 1
    }
    power_state=$(LC_ALL=C pmset -g batt 2>&1 || true)
    thermal_state=$(LC_ALL=C pmset -g therm 2>&1 || true)
    ac_power=false
    case $power_state in
        *"Now drawing from 'AC Power'"*) ac_power=true ;;
    esac
    no_thermal_warning=false
    if printf '%s\n' "$thermal_state" |
           grep -Fq 'No thermal warning level has been recorded' &&
       printf '%s\n' "$thermal_state" |
           grep -Fq 'No performance warning level has been recorded'; then
        no_thermal_warning=true
    elif printf '%s\n' "$thermal_state" |
             grep -Eq 'CPU_Scheduler_Limit[[:space:]]*=[[:space:]]*100' &&
         printf '%s\n' "$thermal_state" |
             grep -Eq 'CPU_Speed_Limit[[:space:]]*=[[:space:]]*100'; then
        no_thermal_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_window_busy_cpu" -v maximum="$max_short_window_busy_cpu" \
        'BEGIN {exit !(actual <= maximum)}' || status=fail
    [ "$ac_power" = true ] || status=fail
    [ "$no_thermal_warning" = true ] || status=fail
    assertions=$(pmset -g assertions 2>&1 || true)
    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 maximum "$max_load_per_cpu" --arg power_state "$power_state" \
        --argjson recent_process_cpu_percent "$recent_process_cpu_percent" \
        --argjson short_window_busy_cpu "$short_window_busy_cpu" \
        --argjson max_short_window_busy_cpu "$max_short_window_busy_cpu" \
        --argjson ac_power "$ac_power" \
        --argjson no_thermal_warning "$no_thermal_warning" \
        --arg thermal_state "$thermal_state" --arg assertions "$assertions" \
        --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: $maximum, power_state: $power_state,
          ac_power: $ac_power,
          recent_aggregate_process_cpu_percent: $recent_process_cpu_percent,
          short_window_busy_cpu_percent: $short_window_busy_cpu,
          max_short_window_busy_cpu_percent: $max_short_window_busy_cpu,
          thermal_state: $thermal_state, no_thermal_warning: $no_thermal_warning,
          assertions: $assertions,
          top_processes: $top_processes}' > "$destination"
    [ "$status" = pass ]
}

for tokens in 4096 65536 131072; do
    position=$((tokens - 32))
    if ! host_receipt preflight "$tokens"; then
        echo "gem: BK128 gate rejected busy host before $tokens-token screen" >&2
        exit 1
    fi
    temporary="$output/seeded-$tokens.json.tmp"
    caffeinate -i "$profile" --shape 31b-full --position "$position" --rows 32 \
        --repetitions "$repetitions" --warmup "$warmup" --seed-prefix \
        --variant compare-bq8-bk128 > "$temporary"
    jq -e --argjson position "$position" --argjson capacity "$tokens" '
      .schema_version == 1 and .engine == "gem-native-c-metal" and
      .shape == "31b-full" and .kv_type == "f32" and
      .position == $position and .rows == 32 and .cache_capacity == $capacity and
      .repetitions == 10 and .warmup == 10 and .seed_prefix == true and
      (.variants["online-bq8"].samples | length) == 10 and
      (.variants["online-bq8-bk128"].samples | length) == 10
    ' "$temporary" >/dev/null
    mv "$temporary" "$output/seeded-$tokens.json"
    if ! host_receipt postflight "$tokens"; then
        echo "gem: BK128 gate rejected changed host after $tokens-token screen" >&2
        exit 1
    fi
done

scripts/analyze-bq8-bk128-gate.sh "$output" "$output/decision.json"
jq . "$output/decision.json"