#!/usr/bin/env bash
# Train, export the policy, and commit every artifact to git so no training
# is ever stranded on one machine. Idempotent and re-runnable: pulls latest
# before training, commits only when there is something new, and a Ctrl-C'd
# run still exports and commits (final.zip is written on the way out).
#
# Usage: ./run_training.sh <experiment> [train.py args...]
# e.g.:  ./run_training.sh run02 --timesteps 20000000 --n-parallel 14 --speedup 16 \
#            --resume checkpoints/run01/final.zip --ent-coef 0.001 --reset-std 0.3
set -uo pipefail

cd "$(dirname "$0")"
EXP="${1:?usage: run_training.sh <experiment> [train.py args...]}"
shift

# Train on the latest code and checkpoints from any machine
git pull --rebase

# Rebuild the exported training binary against the code we just pulled, so a
# --exported-binary run never trains on a stale snapshot. Only when the
# caller has already opted into the exported-binary flow (training/build/
# exists from a prior export) — a source-run caller pays no extra latency.
# Fails the whole run rather than silently falling through to train.py
# against a stale or partially-rebuilt binary (export_linux.sh sets -e, but
# that only exits *it*, not this script).
EXTRA_ARGS=()
if [ -d build ]; then
    ./export_linux.sh || { echo "export_linux.sh failed — aborting" >&2; exit 1; }
    # Standing entry points (next_run.sh, curriculum.py) don't know about the
    # exported binary and never pass --exported-binary themselves; default it
    # here so opting in (by ever running export_linux.sh once) actually gets
    # used, not just kept up to date. An explicit --exported-binary in "$@"
    # still wins (argparse: later occurrence overrides).
    EXTRA_ARGS=(--exported-binary build/CosmicClash.x86_64)
fi

# Let Ctrl-C stop train.py without killing this script, so the export and
# commit below still run
trap ':' INT
.venv/bin/python train.py --experiment "$EXP" "${EXTRA_ARGS[@]}" "$@"
trap - INT

if [ ! -f "checkpoints/$EXP/final.zip" ]; then
    echo "No checkpoints/$EXP/final.zip — nothing to export or commit" >&2
    exit 1
fi

# Export for in-game use (parity-checked); models live in Game/bots/
.venv/bin/python export_policy.py "checkpoints/$EXP/final.zip" "../Game/bots/$EXP.json"

git add -A checkpoints logs eval_history.json "../Game/bots"
if git diff --cached --quiet; then
    echo "Nothing new to commit"
else
    git commit -m "chore(training): Add $EXP checkpoints, logs, and exported policy"
    git push
fi