#!/usr/bin/env bash
# Prepare this bundle for training. Run once, from the bundle root:
#
#     bash scripts/setup.sh [options]
#
# Options:
#   --with-deps        also pip-install runtime dependencies from tcod/ (torch,
#                      vllm, flash-attn, ...). Skip this if the node already has
#                      a working Trinity/TCOD environment -- the default only
#                      installs the vendored tcod package itself, with --no-deps.
#   --models DIR       reuse models already on the node. Symlinks DIR/Qwen3-1.7B
#                      and DIR/Qwen3-32B into ./models instead of downloading.
#   --skip-models      do not download or link models (you will wire ./models up
#                      yourself before running).
#
# What it does, in order:
#   1. sanity-checks python / java
#   2. pip-installs the vendored tcod/ (editable, --no-deps by default)
#   3. copies overlay/trinity/ over tcod/trinity/  -- this is the FutureBridge
#      workflow overlay; see NOTICE
#   4. verifies all five ScienceWorld workflow classes import
#   5. puts Qwen3-1.7B and Qwen3-32B under ./models
#   6. downloads the ScienceWorld task split into ./data/scienceworld
#   7. creates ./outputs/{tmp,ray_tmp,buffers,checkpoints,logs}
#
# Idempotent: every step checks before doing work, so re-running after a partial
# failure is safe.

set -euo pipefail

BUNDLE=$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)
cd "$BUNDLE"

WITH_DEPS=false
SKIP_MODELS=false
EXT_MODELS=""
while [ $# -gt 0 ]; do
    case "$1" in
        --with-deps)   WITH_DEPS=true; shift ;;
        --skip-models) SKIP_MODELS=true; shift ;;
        --models)      EXT_MODELS=${2:-}; shift 2 ;;
        *) echo "unknown option: $1" >&2; exit 2 ;;
    esac
done

log() { printf '\n[setup] %s\n' "$*"; }
die() { printf '[setup] ERROR: %s\n' "$*" >&2; exit 1; }

log "bundle root: $BUNDLE"

# ---------------------------------------------------------------- sanity ----
command -v python >/dev/null 2>&1 || die "no 'python' on PATH"
PYV=$(python -c 'import sys; print("%d.%d" % sys.version_info[:2])')
log "python $PYV"
[ "$PYV" = "3.10" ] || log "WARNING: the release targets python 3.10, found $PYV"

if command -v java >/dev/null 2>&1; then
    log "java: $(java -version 2>&1 | head -1)"
else
    log "WARNING: no 'java' on PATH -- ScienceWorld starts a JVM per episode and"
    log "         will fail at runtime. Needs Java 17+ (e.g. 'module load openjdk')."
fi

# ------------------------------------------------------------ tcod + deps ----
[ -d tcod/trinity ] || die "tcod/ is missing or incomplete -- re-download the bundle"

if [ "$WITH_DEPS" = true ]; then
    log "installing runtime dependencies (--with-deps)"
    # flash-attn compiles against an installed torch, so it cannot come from the
    # same resolution pass; install it last, without build isolation.
    TMP_REQ=$(mktemp)
    grep -v '^flash-attn' requirements.txt > "$TMP_REQ"
    python -m pip install -r "$TMP_REQ"
    rm -f "$TMP_REQ"
    python -m pip install flash-attn==2.8.1 --no-build-isolation || \
        log "WARNING: flash-attn failed to build; training may still work without it"
fi

log "installing the vendored tcod package (editable, --no-deps)"
python -m pip install -e tcod --no-deps

python -c "import scienceworld" 2>/dev/null \
    || { log "installing scienceworld==1.2.2"; python -m pip install 'scienceworld==1.2.2'; }

# ---------------------------------------------------------------- overlay ----
log "applying the FutureBridge workflow overlay onto tcod/"
cp -R overlay/trinity/. tcod/trinity/

log "verifying the five ScienceWorld workflow classes import"
python - <<'EOF'
import importlib
mods = [
    ("trinity.common.workflows.envs.TCOD.scienceworld.OPD_workflow",
     "OnPolicyDistillVerlAgentScienceworldWorkflow"),
    ("trinity.common.workflows.envs.TCOD.scienceworld.guided_opd_workflow",
     "GuidedOPDScienceworldWorkflow"),
    ("trinity.common.workflows.envs.TCOD.scienceworld.TCOD_b2f_workflow",
     "TCOD_b2f_scienceworld_workflow"),
    ("trinity.common.workflows.envs.TCOD.scienceworld.TCOD_f2b_workflow",
     "TCOD_f2b_scienceworld_workflow"),
    ("trinity.common.workflows.envs.TCOD.scienceworld.futurebridge_workflow",
     "FutureBridgeScienceWorldWorkflow"),
]
for mod, cls in mods:
    getattr(importlib.import_module(mod), cls)
    print(f"  OK  {cls}")
EOF

# ----------------------------------------------------------------- models ----
mkdir -p models
if [ "$SKIP_MODELS" = true ]; then
    log "skipping models (--skip-models); wire ./models/Qwen3-1.7B and ./models/Qwen3-32B up yourself"
elif [ -n "$EXT_MODELS" ]; then
    log "linking models from $EXT_MODELS"
    for m in Qwen3-1.7B Qwen3-32B; do
        [ -d "$EXT_MODELS/$m" ] || die "$EXT_MODELS/$m does not exist"
        ln -sfn "$EXT_MODELS/$m" "models/$m"
        echo "  linked models/$m -> $EXT_MODELS/$m"
    done
else
    HF_CLI=$(command -v hf || command -v huggingface-cli || true)
    [ -n "$HF_CLI" ] || die "neither 'hf' nor 'huggingface-cli' found (pip install huggingface_hub), \
or re-run with --models DIR / --skip-models"
    for pair in "Qwen/Qwen3-1.7B:Qwen3-1.7B" "Qwen/Qwen3-32B:Qwen3-32B"; do
        repo=${pair%%:*}; name=${pair##*:}
        if [ -f "models/$name/config.json" ]; then
            log "models/$name already present, skipping"
        else
            log "downloading $repo -> models/$name"
            "$HF_CLI" download "$repo" --local-dir "models/$name"
        fi
    done
fi

# ------------------------------------------------------------------- data ----
log "preparing the ScienceWorld task split"
python scripts/prepare_data.py --out data/scienceworld

# ---------------------------------------------------------------- outputs ----
mkdir -p outputs/tmp outputs/ray_tmp outputs/buffers outputs/checkpoints outputs/logs

log "setup complete."
log "next:  bash scripts/run.sh ftb          (or opd | guided_opd | tcod_b2f | tcod_f2b)"