#!/bin/bash
# Chain: zero-shot baseline eval -> QLoRA training (crash-resume) -> adapter eval.
# Detach with: bash -c './overnight_diffusiongemma.sh & disown'
# All paths overridable via env (override via env).

set -uo pipefail

MODEL="${MODEL:-./diffusiongemma-26B-A4B-it-4bit}"
DATA="${DATA:-./data}"
ADAPTER="${ADAPTER:-./adapters/tool-selector}"
PY="${PY:-python3}"
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
TOTAL_STEPS="${TOTAL_STEPS:-250}"
MAX_RETRIES="${MAX_RETRIES:-6}"
TS=$(date +%Y%m%dT%H%M%S)
LOG="$DATA/chain_${TS}.log"

# smaller Metal command buffers: macOS's GPU interactivity watchdog kills long
# command buffers when the console session is active (kIOGPUCommandBuffer...
# ImpactingInteractivity); short buffers avoid the kill at ~25% throughput cost
export MLX_MAX_OPS_PER_BUFFER=4
export MLX_MAX_MB_PER_BUFFER=20

log() { echo "[$(date +%H:%M:%S)] $*" >> "$LOG"; }
memsnap() { log "memsnap: $(sysctl -n vm.swapusage) | $(memory_pressure 2>/dev/null | tail -1)"; }

latest_checkpoint() {
    ls "$ADAPTER" 2>/dev/null | grep -E '^[0-9]{7}_adapters\.safetensors$' | sort | tail -1
}

mkdir -p "$DATA"
log "=== chain start (steps=$TOTAL_STEPS, buffer caps ops=$MLX_MAX_OPS_PER_BUFFER mb=$MLX_MAX_MB_PER_BUFFER) ==="
memsnap

# Idempotency guard: never silently retrain over a FINISHED adapter; a partial
# best-val adapter with no step checkpoints (crash inside the first save window)
# is archived and training starts fresh
if [ -f "$ADAPTER/adapters.safetensors" ] && [ -z "$(latest_checkpoint)" ]; then
    BEST_STEP=$("$PY" -c "import json; print(json.load(open('$ADAPTER/best.json'))['step'])" 2>/dev/null || echo 0)
    if [ "$BEST_STEP" -ge "$TOTAL_STEPS" ]; then
        log "FATAL: finished adapter present (best step $BEST_STEP >= $TOTAL_STEPS) — refusing to overwrite"
        exit 1
    fi
    log "partial adapter (best step $BEST_STEP, no checkpoints) — archiving to ${ADAPTER}.stale.$TS"
    mv "$ADAPTER" "${ADAPTER}.stale.$TS"
fi

# Phase 1: zero-shot baseline on the clean test split
log "phase 1: zero-shot baseline eval"
caffeinate -ims "$PY" "$SCRIPT_DIR/diffusion_eval.py" \
    --model "$MODEL" --test "$DATA/test.jsonl" \
    --out "$DATA/eval_zeroshot_clean.json" --check-template >> "$LOG" 2>&1
log "phase 1 exit=$?"
memsnap

# Phase 2: training with crash-resume (Metal watchdog kills are intermittent).
# NOTE: bash 3.2 + set -u kills "${ARR[@]}" on EMPTY arrays — use guarded expansion.
TRAIN_EXIT=1
for attempt in $(seq 1 "$MAX_RETRIES"); do
    CKPT=$(latest_checkpoint)
    RESUME_ARGS=()
    if [ -n "$CKPT" ]; then
        STEP=$((10#$(echo "$CKPT" | cut -d_ -f1)))
        if [ "$STEP" -ge "$TOTAL_STEPS" ]; then TRAIN_EXIT=0; log "phase 2: already complete at step $STEP"; break; fi
        RESUME_ARGS=(--resume-file "$ADAPTER/$CKPT" --start-step "$STEP")
        log "phase 2 attempt $attempt: resuming from step $STEP"
    else
        log "phase 2 attempt $attempt: fresh start"
    fi
    caffeinate -ims "$PY" "$SCRIPT_DIR/diffusion_lora_train.py" \
        --model "$MODEL" --data "$DATA" --adapter-path "$ADAPTER" \
        --steps "$TOTAL_STEPS" --val-every 25 --save-every 50 \
        ${RESUME_ARGS[@]+"${RESUME_ARGS[@]}"} >> "$LOG" 2>&1
    TRAIN_EXIT=$?
    log "phase 2 attempt $attempt exit=$TRAIN_EXIT"
    [ "$TRAIN_EXIT" -eq 0 ] && break
    memsnap
    sleep 30
done
memsnap

# Phase 3: adapter eval (best-val adapter)
if [ "$TRAIN_EXIT" -eq 0 ] && [ -f "$ADAPTER/adapters.safetensors" ]; then
    log "phase 3: adapter eval"
    caffeinate -ims "$PY" "$SCRIPT_DIR/diffusion_eval.py" \
        --model "$MODEL" --adapter "$ADAPTER" --test "$DATA/test.jsonl" \
        --out "$DATA/eval_adapter_clean.json" >> "$LOG" 2>&1
    log "phase 3 exit=$?"
else
    log "phase 3 SKIPPED (train exit=$TRAIN_EXIT)"
fi

memsnap
log "=== chain done ==="