#!/usr/bin/env bash
# Run on the named Spark after source, isolated Python, data and artifacts arrive.
# Usage: bash scripts/start_spark_candidate.sh spark-a|spark-b /absolute/fresh/setup-dir
# This script never connects over SSH. Successful setup leaves a detached,
# train-only campaign; its status and stop/resume commands belong to that campaign.
set -euo pipefail

if [[ $# -ne 2 ]]; then
  echo 'Usage: start_spark_candidate.sh spark-a|spark-b /absolute/fresh/setup-dir' >&2
  exit 2
fi
SPARK_NAME="$1"
SETUP_DIR="$2"
case "$SPARK_NAME" in
  spark-a) EXPECTED_HOST=spark-d1b4; MODEL_NAME=Qwen3-4B-Instruct-2507-cdbee75f ;;
  spark-b) EXPECTED_HOST=spark-3e2a; MODEL_NAME=Qwen3.5-2B-15852e8c ;;
  *) echo 'Expected spark-a or spark-b.' >&2; exit 2 ;;
esac
case "$SETUP_DIR" in
  /*) ;;
  *) echo 'SETUP_DIR must be an absolute, fresh directory.' >&2; exit 2 ;;
esac
if [[ -e "$SETUP_DIR" || -L "$SETUP_DIR" ]]; then
  echo "Refusing an existing SETUP_DIR: $SETUP_DIR" >&2
  exit 2
fi
mkdir -p -- "$(dirname -- "$SETUP_DIR")"
mkdir -- "$SETUP_DIR"
SETUP_STAGE=initializing
finish() {
  local setup_exit=$?
  local setup_status=failed
  [[ $setup_exit -eq 0 ]] && setup_status=completed
  set +e
  printf '%s\n' "$setup_exit" > "$SETUP_DIR/exit_code"
  printf '{"stage":"%s","status":"%s","exit_code":%d,"finished_utc":"%s"}\n' \
    "$SETUP_STAGE" "$setup_status" "$setup_exit" "$(date -u +%Y-%m-%dT%H:%M:%SZ)" \
    > "$SETUP_DIR/status.json"
  printf '%s %s exit_code=%s stage=%s\n' "$(date -u +%Y-%m-%dT%H:%M:%SZ)" \
    "$setup_status" "$setup_exit" "$SETUP_STAGE" | tee -a "$SETUP_DIR/stages.log"
}
trap finish EXIT
trap 'exit 130' INT
trap 'exit 143' TERM
stage() {
  SETUP_STAGE="$1"
  printf '%s\n' "$SETUP_STAGE" > "$SETUP_DIR/stage"
  printf '%s %s\n' "$(date -u +%Y-%m-%dT%H:%M:%SZ)" "$SETUP_STAGE" \
    | tee -a "$SETUP_DIR/stages.log"
}

cd "$(dirname "$0")/.."
stage preflight
[[ "$(hostname)" == "$EXPECTED_HOST" ]] || { echo "Expected host $EXPECTED_HOST" >&2; exit 2; }
printf '0\n' > /proc/self/oom_score_adj
export PYTHONUNBUFFERED=1 PYTHONNOUSERSITE=1 TOKENIZERS_PARALLELISM=false
export HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 HF_HUB_DISABLE_XET=1
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
export OPENSYSONE_PYTHON=/home/andy/ai/envs/opensysone/bin/python
export OPENSYSONE_RUN_ROOT=/home/andy/ai/opensysone/runs
RUN_PYTHON="$OPENSYSONE_PYTHON"
MODEL_PATH="/home/andy/ai/models/opensysone/$MODEL_NAME"
MODEL_MANIFEST="/home/andy/ai/opensysone/bootstrap/$MODEL_NAME-files.json"
DATASET=/home/andy/ai/opensysone/data/public-decisions-v1-20260916
[[ -x "$RUN_PYTHON" && -d "$DATASET" && -f "$MODEL_MANIFEST" ]]
git rev-parse HEAD > "$SETUP_DIR/source_commit"
git status --porcelain > "$SETUP_DIR/source_status"

stage model_copy_verification
timeout --signal=TERM --kill-after=30s 5m "$RUN_PYTHON" - \
  "$MODEL_PATH" "$MODEL_MANIFEST" "$SETUP_DIR/model-copy-verified.json" <<'PY'
import hashlib, json, pathlib, sys
from datetime import datetime, timezone
model, manifest_path, output = map(pathlib.Path, sys.argv[1:])
manifest = json.loads(manifest_path.read_text())
if manifest['model'] != model.name or len(manifest['files']) != 13:
    raise ValueError('Unexpected model manifest identity or file count')
verified = {}
for name, expected in manifest['files'].items():
    relative = pathlib.Path(name)
    if relative.is_absolute() or '..' in relative.parts:
        raise ValueError('Model manifest path escapes model directory')
    path = model / relative
    if not path.is_file() or path.stat().st_size != expected['bytes']:
        raise ValueError(f'Missing or incomplete model file: {name}')
    digest = hashlib.sha256()
    with path.open('rb') as handle:
        for chunk in iter(lambda: handle.read(8 * 1024 * 1024), b''):
            digest.update(chunk)
    if digest.hexdigest() != expected['sha256']:
        raise ValueError(f'Model checksum mismatch: {name}')
    verified[name] = expected
proof = {'model': str(model), 'manifest': str(manifest_path),
         'manifest_sha256': hashlib.sha256(manifest_path.read_bytes()).hexdigest(),
         'files': verified, 'verified_utc': datetime.now(timezone.utc).isoformat()}
output.write_text(json.dumps(proof, indent=2) + '\n')
print(json.dumps({'event': 'model_copy_verified', 'model': str(model), 'files': len(verified)}))
PY

memory_preflight() {
  timeout 15s free -b
  timeout 15s nvidia-smi --query-compute-apps=pid,process_name,used_memory --format=csv
}

if [[ "$SPARK_NAME" == spark-a ]]; then
  stage warm_start_pilot
  memory_preflight
  bash scripts/run_experiment.sh train \
    --warm-start /home/andy/ai/opensysone/runs/20260916T183823Z-train/artifacts/checkpoint.pt \
    --model "$MODEL_PATH" --dataset "$DATASET" --steps 8 --epochs 3 \
    --rank 8 --alpha 16 --lr 3e-5 --head-lr 3e-5 --seed 431 \
    --effective-batch 4 --branch-batch-size 1 --max-tokens 512 --two-pass \
    --schedule-steps 7500 --validation-per-family 128 \
    --eval-steps 500 --save-steps 250 --save-seconds 900 --patience 8
  IFS= read -r PILOT_RUN < "$OPENSYSONE_RUN_ROOT/LAST_EXPERIMENT_RUN"
  [[ -f "$PILOT_RUN/exit_code" && "$(cat "$PILOT_RUN/exit_code")" == 0 ]]
  PILOT="$PILOT_RUN/artifacts"
  REFERENCE="$PILOT/validation_step_000008_predictions.json"
else
  PILOT=/home/andy/ai/opensysone/runs/20260916T182352Z-train/artifacts
  REFERENCE="$PILOT/validation_step_000040_predictions.json"
fi
printf '%s\n' "$PILOT" > "$SETUP_DIR/pilot_path"
[[ -f "$PILOT/checkpoint.pt" && -f "$PILOT/best.pt" && -f "$REFERENCE" ]]

stage artifact_verification
memory_preflight
timeout --signal=TERM --kill-after=30s 15m "$RUN_PYTHON" scripts/verify_artifact.py \
  --checkpoint "$PILOT/checkpoint.pt" --reference "$REFERENCE" \
  --inference-max-tokens 1024 --output "$SETUP_DIR/verify"

stage campaign_launch
memory_preflight
timeout --signal=TERM --kill-after=30s 2m "$RUN_PYTHON" scripts/launch_24h.py \
  --pilot "$PILOT" --train-only --training-deadline 2026-09-17T16:00:00Z \
  --deadline 2026-09-17T18:16:10Z --inference-max-tokens 1024 \
  | tee "$SETUP_DIR/launcher.json"
cp -- "$OPENSYSONE_RUN_ROOT/LAST_CAMPAIGN" "$SETUP_DIR/campaign_path"
"$RUN_PYTHON" - "$SETUP_DIR" <<'PY'
import json, pathlib, sys
setup = pathlib.Path(sys.argv[1])
launched = json.loads((setup / 'launcher.json').read_text())
if launched['campaign'] != (setup / 'campaign_path').read_text().strip():
    raise RuntimeError('Campaign pointer differs from launcher result')
print(json.dumps({'event': 'train_only_campaign_launched', **launched}))
PY
stage campaign_running