#!/usr/bin/env bash
#
# Everything that must hold before run 1 starts.
#
# Each of these exists because something in it was actually broken. The
# gradient accumulation held every micro batch live, resume died on the first
# optimizer step, sample.py had never met a live MLX install, and the MTP index
# arithmetic is the bug GOAL predicted would be the expensive one. Running the
# suite takes a few minutes. Discovering any of it on day four costs days.
#
# Usage:
#   ./scripts/preflight.sh                    # full gate, audits data/shards
#   SKIP_CORPUS=1 ./scripts/preflight.sh      # code only, before the corpus exists
#
set -u

PY=${PY:-.venv/bin/python}
CONFIG=${CONFIG:-config/run1.json}
SKIP_CORPUS=${SKIP_CORPUS:-0}

pass=0
fail=0
declare -a FAILED

check() {
  local name=$1; shift
  printf '%-42s' "$name"
  if "$@" >"/tmp/preflight_$(echo "$name" | tr -c 'a-zA-Z0-9' '_').log" 2>&1; then
    echo "PASS"
    pass=$((pass + 1))
  else
    echo "FAIL"
    fail=$((fail + 1))
    FAILED+=("$name")
  fi
}

echo "=== wisp preflight ==="
echo

check "python files compile" $PY -m py_compile \
  model.py data.py checkpoint_fs.py train.py sample.py \
  scripts/test_data_contract.py \
  scripts/test_checkpoint_transaction.py \
  scripts/test_corpus_quality.py \
  scripts/training_data_contract.py \
  scripts/test_training_data_contract.py \
  scripts/derive_no_fim.py scripts/test_derive_no_fim.py \
  scripts/test_sampler_resume.py \
  scripts/test_prepare_data.py \
  scripts/e2_contract.py scripts/test_e2_contract.py \
  scripts/train_tokenizer.py scripts/prepare_data.py scripts/corpus.py \
  scripts/benchmark_tokenizer_vocab.py \
  scripts/build_eval_holdout.py scripts/check_holdout_overlap.py \
  scripts/check_eval_holdout.py \
  scripts/eval_acceptance.py scripts/eval_pairs.py \
  scripts/compare_acceptance.py scripts/compare_format_ablation.py \
  scripts/eval_rollout.py \
  scripts/rollout_metrics.py scripts/eval_validation.py \
  scripts/validation_metrics.py scripts/release_audit.py \
  scripts/export_hf.py \
  scripts/export_quantized.py scripts/test_export_quantized.py \
  scripts/publish_hf.py scripts/test_publish_hf.py \
  scripts/mutation_audit.py \
  scripts/hf_metadata.py scripts/summarize_training.py \
  scripts/supervisor_state.py scripts/supervisor_heartbeat.py \
  scripts/supervisor_fixture.py \
  scripts/bench.py scripts/audit_corpus.py scripts/test_eval_pairs.py \
  scripts/test_compare_acceptance.py \
  scripts/test_compare_format_ablation.py scripts/test_rollout_metrics.py \
  scripts/test_rollout_replay.py scripts/test_mtp_trace.py \
  scripts/verify_rollout_replay.py scripts/test_verify_rollout_replay.py \
  scripts/test_release_rollout_v3.py \
  scripts/test_validation_metrics.py \
  scripts/test_release_audit.py \
  scripts/test_build_eval_holdout.py scripts/test_holdout_overlap.py \
  scripts/test_tokenizer_vocab.py scripts/test_export_metadata.py \
  scripts/test_summarize_training.py

check "holdout builder selection" $PY scripts/test_build_eval_holdout.py
check "eval pair construction" $PY scripts/test_eval_pairs.py
check "acceptance control comparison" $PY scripts/test_compare_acceptance.py
check "format-ablation comparison" $PY scripts/test_compare_format_ablation.py
check "adaptive rollout metrics" $PY scripts/test_rollout_metrics.py
check "rollout branch replay" $PY scripts/test_rollout_replay.py
check "MTP generation trace" $PY scripts/test_mtp_trace.py
check "independent rollout replay" $PY scripts/test_verify_rollout_replay.py
check "release rollout v3 audit" $PY scripts/test_release_rollout_v3.py
check "final validation metrics" $PY scripts/test_validation_metrics.py
check "final release audit" $PY scripts/test_release_audit.py
check "training data contract" $PY scripts/test_data_contract.py
check "atomic checkpoint transaction" $PY scripts/test_checkpoint_transaction.py
check "Hub corpus quality schema" $PY scripts/test_corpus_quality.py
check "training-data disclosure" $PY scripts/test_training_data_contract.py
check "deterministic no-FIM derivation" $PY scripts/test_derive_no_fim.py
check "training sampler resume" $PY scripts/test_sampler_resume.py
check "corpus build setup" $PY scripts/test_prepare_data.py
check "E2 format-ablation contract" $PY scripts/test_e2_contract.py
check "supervisor lifecycle" bash scripts/test_supervisor.sh
check "holdout overlap scanner" $PY scripts/test_holdout_overlap.py
check "tokenizer merge-prefix truncation" $PY scripts/test_tokenizer_vocab.py
check "Hugging Face export metadata" $PY scripts/test_export_metadata.py
check "rollback-aware training telemetry" $PY scripts/test_summarize_training.py
check "MTP index alignment" $PY scripts/test_mtp_indices.py
check "chunked cross entropy equivalence" $PY scripts/test_ce_chunk.py
check "native GQA attention equivalence" $PY scripts/test_gqa_attention.py
check "quantized-export comparison math" $PY scripts/test_export_quantized.py
check "Hugging Face publication gate" $PY scripts/test_publish_hf.py
check "checkpoint resume" $PY scripts/test_resume.py

if [ "$SKIP_CORPUS" = "0" ]; then
  IDX=$($PY -c "import json;print(json.load(open('$CONFIG'))['data_index'])")
  TOK=$($PY -c "import json;print(json.load(open('$CONFIG'))['tokenizer_path'])")
  check "corpus audit" $PY scripts/audit_corpus.py --index "$IDX" --tokenizer "$TOK"
else
  echo "corpus audit                               SKIPPED"
fi

echo
echo "=== $pass passed, $fail failed ==="
if [ $fail -gt 0 ]; then
  echo
  echo "failures:"
  for f in "${FAILED[@]}"; do
    echo "  - $f  (see /tmp/preflight_$(echo "$f" | tr -c 'a-zA-Z0-9' '_').log)"
  done
  echo
  echo "DO NOT LAUNCH. Fix these first."
  exit 1
fi

echo
echo "cleared for launch:"
echo "  caffeinate -dims ./scripts/supervise.sh $CONFIG"
exit 0