#!/bin/sh
# Capture a production-route long-prefill stage profile before selecting a kernel axis.
set -eu

if [ "$#" -ne 2 ]; then
    echo "usage: $0 MODEL.gguf OUTPUT_DIR" >&2
    exit 2
fi
model=$1
output=$2
build_dir=${GEM_PREFILL_STAGE_BUILD:-build}

[ -r "$model" ] || { echo "gem: model is not readable: $model" >&2; exit 2; }
[ ! -e "$output" ] || { echo "gem: output already exists: $output" >&2; exit 2; }
if command -v pgrep >/dev/null 2>&1 && pgrep -f '[/]gem serve' >/dev/null 2>&1; then
    echo "gem: refusing to overlap a running native server Metal workload" >&2
    exit 1
fi
for binary in "$build_dir/gem" "$build_dir/gem-long-prefill-stage-profile" \
              "$build_dir/gem-metal-counter-info"; do
    [ -x "$binary" ] || { echo "gem: required binary is missing: $binary" >&2; exit 2; }
done

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
    }
fi
actual_model_sha256=$(shasum -a 256 "$model" | awk '{print $1}')
if [ -n "$model_sha256" ] && [ "$actual_model_sha256" != "$model_sha256" ]; then
    echo "gem: model digest does not match GEM_MODEL_SHA256" >&2
    exit 1
fi
model_sha256=$actual_model_sha256

mkdir -p "$output"
date -u '+%Y-%m-%dT%H:%M:%SZ' > "$output/run-start-utc.txt"
printf '%s  %s\n' "$model_sha256" "$model" > "$output/model-sha256.txt"
sw_vers > "$output/macos-version.txt"
shasum -a 256 "$build_dir/gem" "$build_dir/gem-long-prefill-stage-profile" \
    "$build_dir/gem-metal-counter-info" > "$output/binary-sha256.txt"
otool -L "$build_dir/gem" "$build_dir/gem-long-prefill-stage-profile" \
    "$build_dir/gem-metal-counter-info" > "$output/native-linkage.txt"
if grep -Eiq 'lib(llama|ggml)' "$output/native-linkage.txt"; then
    echo "gem: production profiler unexpectedly links llama/ggml" >&2
    exit 1
fi

"$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 .kv_bytes == 23194501120 and
  (.attention_kernel | contains("31b-full-online-row2")) and
  (.attention_kernel | contains("full-decode-value-dim32-split16@8192"))
' "$output/print-config.json" >/dev/null

"$build_dir/gem-metal-counter-info" > "$output/metal-counter-info.json"
jq -e '
  .schema_version == 1 and .workload == "metal-counter-info" and
  .device.has_unified_memory == true and .device.gpu_families.apple9 == true and
  .device.max_threadgroup_memory_bytes >= 32768 and
  .sampling_points.stage_boundary == true and
  (.counter_sets | type == "array" and length > 0)
' "$output/metal-counter-info.json" >/dev/null
scripts/analyze-metal-counter-info.sh "$output/metal-counter-info.json" \
    "$output/metal-counter-analysis.json"
jq -e '
  .instrumentation.stage_timestamps_available == true and
  (.instrumentation.selected == "public-dispatch-counter-sampling" or
   .instrumentation.selected == "stage-timestamps-plus-xcode-gpu-trace")
' "$output/metal-counter-analysis.json" >/dev/null

"$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"

jq -e '
  .screening.promotion_eligible == false and
  .source.layer_count == 60 and .source.sliding_layer_count == 50 and
  .source.full_layer_count == 10 and .source.repetitions == 3 and
  .source.prefix_mode == "seeded-kv" and .source.prefix_tokens == 131040 and
  .source.tile_tokens == 32
' "$output/prefill-stage-analysis-131040.json" >/dev/null
date -u '+%Y-%m-%dT%H:%M:%SZ' > "$output/run-end-utc.txt"
jq . "$output/prefill-stage-analysis-131040.json"