#!/usr/bin/env bash
set -euo pipefail

if [[ $# -lt 2 || $# -gt 4 ]]; then
  echo "usage: $0 <uniform|shuffle|residual> <run-id> [proposal.json] [config.yaml]" >&2
  exit 64
fi

arm=$1
run_id=$2
proposal=${3:-${OPTR_TEMPORAL_PROPOSAL_PATH:-}}
project="${OPTR_PROJECT_ROOT:-$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)}"
python="${OPTR_PYTHON:-python}"
config=${4:-"${project}/configs/florence_xtrainer2_pack_phone_residual_20ep_to_60ep.yaml"}
accelerate_config="${OPTR_ACCELERATE_CONFIG:-${project}/configs/deepspeed_florence_8gpu.yaml}"
num_processes="${OPTR_NUM_PROCESSES:-8}"

missing=()
for var in OPTR_RUN_ROOT FLORENCE_MODEL_PATH XTRAINER2_DATA_ROOT \
  XTRAINER2_PACK_PHONE_PROVENANCE_PATH XTRAINER2_PACK_PHONE_FINGERPRINT_PATH \
  XTRAINER2_PACK_PHONE_MANIFEST_PATH XTRAINER2_PACK_PHONE_RESUME_DIR \
  XTRAINER2_PACK_PHONE_RESUME_INTEGRITY_PATH; do
  [[ -n "${!var:-}" ]] || missing+=("$var")
done
if (( ${#missing[@]} )); then
  echo "error: required environment variables are unset: ${missing[*]}" >&2
  exit 78
fi

[[ "$num_processes" == 8 ]] || {
  echo "error: OPTR continuation requires exactly 8 processes (got $num_processes)" >&2
  exit 78
}
[[ "${FLORENCE_MODEL_PATH,,}" == *florence* && -d "$FLORENCE_MODEL_PATH" ]] || {
  echo "error: FLORENCE_MODEL_PATH must be an existing Florence directory" >&2
  exit 66
}
[[ -d "$XTRAINER2_DATA_ROOT" ]] || {
  echo "error: XTRAINER2_DATA_ROOT is not a directory: $XTRAINER2_DATA_ROOT" >&2
  exit 66
}
for path in "$config" "$accelerate_config" \
  "$XTRAINER2_PACK_PHONE_PROVENANCE_PATH" \
  "$XTRAINER2_PACK_PHONE_FINGERPRINT_PATH" \
  "$XTRAINER2_PACK_PHONE_MANIFEST_PATH" \
  "$XTRAINER2_PACK_PHONE_RESUME_INTEGRITY_PATH"; do
  [[ -f "$path" ]] || { echo "error: required file missing: $path" >&2; exit 66; }
done
[[ -d "$XTRAINER2_PACK_PHONE_RESUME_DIR" ]] || {
  echo "error: resume checkpoint directory missing: $XTRAINER2_PACK_PHONE_RESUME_DIR" >&2
  exit 66
}
[[ "$(basename "$XTRAINER2_PACK_PHONE_RESUME_DIR")" == steps_17460_accelerate ]] || {
  echo "error: exact common-start must be steps_17460_accelerate" >&2
  exit 78
}
if ! command -v "$python" >/dev/null 2>&1 && [[ ! -x "$python" ]]; then
  echo "error: python interpreter not found: $python" >&2
  exit 69
fi

case "$arm" in
  uniform)
    [[ -z "$proposal" ]] || {
      echo "error: uniform arm does not accept a proposal" >&2
      exit 64
    }
    proposal_args=(--datasets.vla_data.temporal_proposal_path=null)
    ;;
  shuffle|residual)
    [[ -n "$proposal" && -f "$proposal" ]] || {
      echo "error: $arm requires a proposal file via CLI or OPTR_TEMPORAL_PROPOSAL_PATH" >&2
      exit 66
    }
    proposal_args=(--datasets.vla_data.temporal_proposal_path="$proposal")
    ;;
  *)
    echo "error: unknown arm: $arm" >&2
    exit 64
    ;;
esac

output_dir="${OPTR_RUN_ROOT}/${run_id}"
[[ ! -e "$output_dir" ]] || {
  echo "error: refusing to reuse run directory: $output_dir" >&2
  exit 73
}

export PYTHONPATH="${project}:${PYTHONPATH:-}"
"$python" - "$XTRAINER2_PACK_PHONE_RESUME_DIR" \
  "$XTRAINER2_PACK_PHONE_RESUME_INTEGRITY_PATH" \
  "$XTRAINER2_PACK_PHONE_MANIFEST_PATH" "${proposal:-}" "$arm" <<'PY'
import json
import sys
from pathlib import Path

checkpoint, inventory_path, manifest_path, proposal_path = map(Path, sys.argv[1:5])
arm = sys.argv[5]
commit_path = checkpoint / "checkpoint.json"
try:
    commit = json.loads(commit_path.read_text(encoding="utf-8"))
except (OSError, json.JSONDecodeError) as exc:
    raise SystemExit(f"error: invalid checkpoint commit {commit_path}: {exc}")
if commit.get("schema_version") != 1:
    raise SystemExit("error: unsupported checkpoint schema")
if int(commit.get("completed_steps", -1)) != 17460:
    raise SystemExit("error: checkpoint is not global step 17460")
if int(commit.get("world_size", -1)) != 8:
    raise SystemExit("error: checkpoint world size is not 8")
files = commit.get("files")
if not isinstance(files, list) or not files:
    raise SystemExit("error: checkpoint commit has no file inventory")
for entry in files:
    path = checkpoint / entry["name"]
    if not path.is_file() or path.stat().st_size != int(entry["bytes"]):
        raise SystemExit(f"error: incomplete checkpoint file: {path}")

try:
    inventory = json.loads(inventory_path.read_text(encoding="utf-8"))
except (OSError, json.JSONDecodeError) as exc:
    raise SystemExit(f"error: invalid checkpoint integrity manifest: {exc}")
expected = [(x["name"], int(x["bytes"])) for x in files]
if inventory.get("schema_version") != 1:
    raise SystemExit("error: checkpoint integrity manifest must use schema 1")
if inventory.get("status") == "locked":
    observed = [(x["name"], int(x["bytes"])) for x in inventory.get("files", [])]
    if observed != expected:
        raise SystemExit("error: checkpoint integrity manifest differs from checkpoint.json")
elif inventory.get("status") == "pass":
    if (
        inventory.get("errors") != []
        or int(inventory.get("completed_steps", -1)) != 17460
        or int(inventory.get("world_size", -1)) != 8
        or int(inventory.get("inventory_files", -1)) != len(files)
    ):
        raise SystemExit("error: common-start preflight report is inconsistent")
else:
    raise SystemExit("error: checkpoint integrity manifest did not pass")

try:
    from starVLA.dataloader.infoflow_sampler import ValidAnchorManifest
    from starVLA.dataloader.temporal_proposal import TemporalProposal
    manifest = ValidAnchorManifest.load(manifest_path)
    if arm != "uniform":
        proposal = TemporalProposal.load(proposal_path, manifest)
        if proposal.score_kind != arm:
            raise ValueError(
                f"{arm} arm received {proposal.score_kind!r} proposal"
            )
except Exception as exc:
    raise SystemExit(f"error: manifest/proposal validation failed: {exc}")
PY

if [[ "${OPTR_PREFLIGHT_ONLY:-0}" == 1 ]]; then
  echo "preflight passed: $arm -> $output_dir"
  exit 0
fi

cd "$project"
export WANDB_MODE="${WANDB_MODE:-offline}"
export PYTHONUNBUFFERED=1
exec "$python" -m accelerate.commands.launch \
  --config_file "$accelerate_config" \
  --num_processes 8 \
  starVLA/training/train_starvla.py \
  --config_yaml "$config" \
  --run_id "$run_id" \
  --run_root_dir "$OPTR_RUN_ROOT" \
  "${proposal_args[@]}"