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

# Prepare the complete LeRobot-v2 T-Rex dataset for VLTA/T-Rex Track-Force
# training. The source dataset is preserved: by default a hard-link clone is
# created, so RGB/tactile videos do not consume a second copy of their blocks.
#
# The Python builder then:
#   1. extracts canonical 250-point SAM2 + CoTracker tracks;
#   2. adds track, visibility, and absolute EEF62 parquet columns;
#   3. validates the existing 60-D (10 x 6) tactile-force column;
#   4. writes force/track/action statistics and the training metadata contract;
#   5. validates every episode before optionally starting training.

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

SOURCE_ROOT="${SOURCE_TREX_DATA_ROOT:-$REPO_ROOT/data/trex_full}"
DATASET_ROOT="${TREX_DATA_ROOT:-$REPO_ROOT/data/trex_full_force}"
TRACK_CACHE="${TRACK_CACHE:-}"
DEFAULT_PYTHON="$REPO_ROOT/../miniconda3/envs/dreamzero/bin/python"
if [[ ! -x "$DEFAULT_PYTHON" ]]; then
    DEFAULT_PYTHON="python"
fi
PYTHON_BIN="${PYTHON_BIN:-$DEFAULT_PYTHON}"
GPU_LIST="${GPUS:-6,7}"
EPISODES_PER_JOB="${EPISODES_PER_JOB:-32}"
GPU_MIN_FREE_MIB="${GPU_MIN_FREE_MIB:-78000}"
GPU_RETRY_FREE_MIB="${GPU_RETRY_FREE_MIB:-78000}"
GPU_POLL_SECONDS="${GPU_POLL_SECONDS:-60}"
GPU_REQUIRE_IDLE="${GPU_REQUIRE_IDLE:-6,7}"
START_EPISODE="${START_EPISODE:-auto}"

OPENPI_ROOT="${OPENPI_ROOT:-/scratch2/home/zhicao/openpi}"
SAM2_LIBS="${SAM2_LIBS:-/scratch1/home/zhicao/physctrl/libs}"
SAM2_MODEL="${SAM2_MODEL:-facebook/sam2-hiera-large}"
CALIB_PATH="${CALIB_PATH:-$REPO_ROOT/assets/trex_camera_calib.json}"
COTRACKER_CHECKPOINT="${COTRACKER_CHECKPOINT:-}"
T_REX_ROOT="${T_REX_ROOT:-/scratch1/home/zhicao/T-Rex}"

RUN_TRAINING=0
VALIDATE_ONLY=0
FORCE_REBUILD=0
ALLOW_IN_PLACE=0
VERIFY_FK="${VERIFY_FK:-0}"
TRAIN_ARGS=()

usage() {
    cat <<EOF
Usage:
  $(basename "$0") [options] [-- training_hydra_overrides...]

Prepare all T-Rex v2 episodes for:
  $REPO_ROOT/scripts/train/trex_track_force_training_wan22.sh

Options:
  --source-root PATH    Original LeRobot-v2 dataset
                        (default: $SOURCE_ROOT)
  --dataset-root PATH   Prepared output dataset
                        (default: $DATASET_ROOT)
  --track-cache PATH    Resumable per-episode track cache
                        (default: DATASET_ROOT/tracks_trex_track_force_v2)
  --gpus LIST           Comma-separated physical GPU IDs/UUIDs
                        (default: $GPU_LIST)
  --start-episode N     First episode to inspect, or 'auto' for the first
                        missing/invalid track cache (default: $START_EPISODE)
  --episodes-per-job N  Episodes per dynamic queue task
                        (default: $EPISODES_PER_JOB)
  --min-free-gpu-mib N  Free GPU memory required before starting a task
                        (default: $GPU_MIN_FREE_MIB)
  --require-idle-gpus L Comma-separated GPUs that must have no other compute
                        process before joining the queue (default: $GPU_REQUIRE_IDLE)
  --python PATH         Python executable (default: $PYTHON_BIN)
  --force               Rebuild already-valid converted episodes
  --verify-fk           Recompute FK during validation (much slower)
  --validate-only       Only validate an already-prepared output dataset
  --in-place            Permit SOURCE_ROOT == DATASET_ROOT
  --train               Start Track-Force training after validation
  -h, --help            Show this help

Examples:
  # Resume from the first missing track on the default exclusive GPUs 6 and 7.
  $(basename "$0")

  # Override automatic resume when a specific restart point is required.
  $(basename "$0") --start-episode 1717

  # Validate only, without loading SAM2 or CoTracker.
  $(basename "$0") --validate-only

  # Prepare and immediately train, forwarding Hydra overrides.
  $(basename "$0") --gpus 0,1,2,3 --train -- max_steps=100000

Environment equivalents:
  SOURCE_TREX_DATA_ROOT, TREX_DATA_ROOT, TRACK_CACHE, GPUS, PYTHON_BIN,
  START_EPISODE, EPISODES_PER_JOB, GPU_MIN_FREE_MIB, GPU_RETRY_FREE_MIB,
  GPU_POLL_SECONDS, GPU_REQUIRE_IDLE, OPENPI_ROOT, SAM2_LIBS, SAM2_MODEL, CALIB_PATH,
  COTRACKER_CHECKPOINT.
EOF
}

need_value() {
    if [[ $# -lt 2 || -z "$2" ]]; then
        echo "ERROR: $1 requires a value" >&2
        exit 2
    fi
}

while [[ $# -gt 0 ]]; do
    case "$1" in
        --source-root)
            need_value "$@"
            SOURCE_ROOT="$2"
            shift 2
            ;;
        --dataset-root)
            need_value "$@"
            DATASET_ROOT="$2"
            shift 2
            ;;
        --track-cache)
            need_value "$@"
            TRACK_CACHE="$2"
            shift 2
            ;;
        --gpus)
            need_value "$@"
            GPU_LIST="$2"
            shift 2
            ;;
        --start-episode)
            need_value "$@"
            START_EPISODE="$2"
            shift 2
            ;;
        --episodes-per-job)
            need_value "$@"
            EPISODES_PER_JOB="$2"
            shift 2
            ;;
        --min-free-gpu-mib)
            need_value "$@"
            GPU_MIN_FREE_MIB="$2"
            shift 2
            ;;
        --require-idle-gpus)
            need_value "$@"
            GPU_REQUIRE_IDLE="$2"
            shift 2
            ;;
        --python)
            need_value "$@"
            PYTHON_BIN="$2"
            shift 2
            ;;
        --force)
            FORCE_REBUILD=1
            shift
            ;;
        --verify-fk)
            VERIFY_FK=1
            shift
            ;;
        --validate-only)
            VALIDATE_ONLY=1
            shift
            ;;
        --in-place)
            ALLOW_IN_PLACE=1
            shift
            ;;
        --train)
            RUN_TRAINING=1
            shift
            ;;
        --)
            shift
            TRAIN_ARGS=("$@")
            break
            ;;
        -h|--help)
            usage
            exit 0
            ;;
        *)
            echo "ERROR: unknown option: $1" >&2
            usage >&2
            exit 2
            ;;
    esac
done

if [[ "$START_EPISODE" != "auto" && ! "$START_EPISODE" =~ ^[0-9]+$ ]]; then
    echo "ERROR: START_EPISODE must be 'auto' or a non-negative integer" >&2
    exit 2
fi

for numeric_setting in \
    "EPISODES_PER_JOB=$EPISODES_PER_JOB" \
    "GPU_MIN_FREE_MIB=$GPU_MIN_FREE_MIB" \
    "GPU_RETRY_FREE_MIB=$GPU_RETRY_FREE_MIB" \
    "GPU_POLL_SECONDS=$GPU_POLL_SECONDS"; do
    setting_name="${numeric_setting%%=*}"
    setting_value="${numeric_setting#*=}"
    if [[ ! "$setting_value" =~ ^[1-9][0-9]*$ ]]; then
        echo "ERROR: $setting_name must be a positive integer, got: $setting_value" >&2
        exit 2
    fi
done

if ! command -v "$PYTHON_BIN" >/dev/null 2>&1; then
    echo "ERROR: Python executable not found: $PYTHON_BIN" >&2
    exit 1
fi

canonical_path() {
    "$PYTHON_BIN" - "$1" <<'PY'
import sys
from pathlib import Path

print(Path(sys.argv[1]).expanduser().resolve())
PY
}

SOURCE_ROOT="$(canonical_path "$SOURCE_ROOT")"
DATASET_ROOT="$(canonical_path "$DATASET_ROOT")"
if [[ -z "$TRACK_CACHE" ]]; then
    TRACK_CACHE="$DATASET_ROOT/tracks_trex_track_force_v2"
fi
TRACK_CACHE="$(canonical_path "$TRACK_CACHE")"
OPENPI_ROOT="$(canonical_path "$OPENPI_ROOT")"
SAM2_LIBS="$(canonical_path "$SAM2_LIBS")"
CALIB_PATH="$(canonical_path "$CALIB_PATH")"
T_REX_ROOT="$(canonical_path "$T_REX_ROOT")"
if [[ -z "$COTRACKER_CHECKPOINT" ]]; then
    COTRACKER_CHECKPOINT="$OPENPI_ROOT/co-tracker/checkpoints/scaled_offline.pth"
else
    COTRACKER_CHECKPOINT="$(canonical_path "$COTRACKER_CHECKPOINT")"
fi

BUILDER="$REPO_ROOT/scripts/data/build_trex_track_force_v2.py"
TRAIN_SCRIPT="$REPO_ROOT/scripts/train/trex_track_force_training_wan22.sh"
CANONICAL_MANIFEST="$DATASET_ROOT/meta/trex_track_force_manifest.json"

for required_file in "$BUILDER" "$TRAIN_SCRIPT"; do
    if [[ ! -f "$required_file" ]]; then
        echo "ERROR: required script is missing: $required_file" >&2
        exit 1
    fi
done

if [[ "$SOURCE_ROOT" == "$DATASET_ROOT" && "$ALLOW_IN_PLACE" != "1" ]]; then
    echo "ERROR: refusing to modify the source dataset in place." >&2
    echo "       Use a different --dataset-root or explicitly pass --in-place." >&2
    exit 1
fi

if [[ "$VALIDATE_ONLY" == "1" ]]; then
    if [[ ! -f "$DATASET_ROOT/meta/info.json" ]]; then
        echo "ERROR: prepared dataset does not exist: $DATASET_ROOT" >&2
        exit 1
    fi
else
    if [[ ! -f "$SOURCE_ROOT/meta/info.json" ]]; then
        echo "ERROR: source is not a LeRobot-v2 dataset: $SOURCE_ROOT" >&2
        exit 1
    fi

    if [[ ! -e "$DATASET_ROOT" ]]; then
        mkdir -p "$(dirname "$DATASET_ROOT")"
        echo "Creating hard-link clone (source remains unchanged):"
        echo "  source: $SOURCE_ROOT"
        echo "  output: $DATASET_ROOT"
        if ! cp -al -- "$SOURCE_ROOT" "$DATASET_ROOT"; then
            echo "ERROR: hard-link clone failed." >&2
            echo "       Source and output must be on the same filesystem." >&2
            echo "       Remove the incomplete output or provide an existing copied dataset." >&2
            exit 1
        fi
    elif [[ ! -f "$DATASET_ROOT/meta/info.json" ]]; then
        echo "ERROR: output exists but is not a resumable dataset: $DATASET_ROOT" >&2
        exit 1
    fi
fi

# Check the complete disk contract before loading expensive tracking models.
TOTAL_EPISODES="$(
    "$PYTHON_BIN" - "$SOURCE_ROOT" "$DATASET_ROOT" <<'PY'
import json
import sys
from pathlib import Path

source = Path(sys.argv[1])
output = Path(sys.argv[2])
output_info = json.loads((output / "meta" / "info.json").read_text())
source_info_path = source / "meta" / "info.json"
source_info = (
    json.loads(source_info_path.read_text())
    if source_info_path.is_file()
    else output_info
)

if source_info.get("codebase_version") != "v2.1":
    raise SystemExit(
        f"source codebase_version must be v2.1, got {source_info.get('codebase_version')!r}"
    )
if output_info.get("total_episodes") != source_info.get("total_episodes"):
    raise SystemExit("source/output total_episodes mismatch")

required_metadata = (
    "modality.json",
    "stats.json",
    "episodes.jsonl",
    "tasks.jsonl",
    "embodiment.json",
)
missing_metadata = [
    name for name in required_metadata if not (output / "meta" / name).is_file()
]
if missing_metadata:
    raise SystemExit(f"dataset is missing metadata files: {missing_metadata}")

features = output_info.get("features", {})
required_shapes = {
    "observation.state": [58],
    "action": [58],
    "observation.tactile_force": [60],
}
for name, shape in required_shapes.items():
    actual = features.get(name, {}).get("shape")
    if actual != shape:
        raise SystemExit(f"{name} shape must be {shape}, got {actual}")

video_keys = (
    "observation.images.head_left",
    "observation.images.left_wrist",
    "observation.images.right_wrist",
)
for key in video_keys:
    if features.get(key, {}).get("dtype") != "video":
        raise SystemExit(f"missing RGB video feature {key}")

total = int(output_info["total_episodes"])
chunk_size = int(output_info.get("chunks_size", 1000))
data_pattern = output_info.get(
    "data_path",
    "data/chunk-{episode_chunk:03d}/episode_{episode_index:06d}.parquet",
)
video_pattern = output_info.get(
    "video_path",
    "videos/chunk-{episode_chunk:03d}/{video_key}/episode_{episode_index:06d}.mp4",
)
missing = []
for episode in range(total):
    values = {"episode_chunk": episode // chunk_size, "episode_index": episode}
    parquet = output / data_pattern.format(**values)
    if not parquet.is_file():
        missing.append(str(parquet))
    for key in video_keys:
        video = output / video_pattern.format(video_key=key, **values)
        if not video.is_file():
            missing.append(str(video))
    if len(missing) >= 20:
        break
if missing:
    preview = "\n".join(f"  - {path}" for path in missing)
    raise SystemExit(f"dataset is incomplete; missing required files:\n{preview}")

print(total)
PY
)"

if [[ ! "$TOTAL_EPISODES" =~ ^[1-9][0-9]*$ ]]; then
    echo "ERROR: invalid episode count: $TOTAL_EPISODES" >&2
    exit 1
fi

FK_FLAG="--no-verify-fk"
if [[ "$VERIFY_FK" == "1" ]]; then
    FK_FLAG="--verify-fk"
fi

if [[ "$VALIDATE_ONLY" == "1" ]]; then
    echo "Validating $TOTAL_EPISODES prepared episodes..."
    "$PYTHON_BIN" "$BUILDER" \
        --dataset-root "$DATASET_ROOT" \
        --track-cache "$TRACK_CACHE" \
        --manifest-path "$CANONICAL_MANIFEST" \
        --all \
        --validate-only \
        "$FK_FLAG"
    echo "Dataset is ready for training: $DATASET_ROOT"
else
    for required_path in \
        "$OPENPI_ROOT/droid" \
        "$COTRACKER_CHECKPOINT" \
        "$SAM2_LIBS/sam2" \
        "$CALIB_PATH" \
        "$T_REX_ROOT/utils/lerobot_common.py"; do
        if [[ ! -e "$required_path" ]]; then
            echo "ERROR: preprocessing dependency is missing: $required_path" >&2
            exit 1
        fi
    done

    "$PYTHON_BIN" - "$SAM2_LIBS" "$OPENPI_ROOT" <<'PY'
import sys
from pathlib import Path

missing = []
for module in ("cv2", "hydra", "numpy", "pandas", "pinocchio", "pyarrow", "torch"):
    try:
        __import__(module)
    except Exception as exc:  # noqa: BLE001
        missing.append(f"{module}: {exc}")
if missing:
    raise SystemExit("missing Python dependencies:\n  " + "\n  ".join(missing))

sam2_libs = Path(sys.argv[1])
openpi_root = Path(sys.argv[2])
sys.path.insert(0, str(sam2_libs))
sys.path.insert(0, str(openpi_root / "droid"))
try:
    from sam2.sam2_image_predictor import SAM2ImagePredictor  # noqa: F401
except Exception as exc:  # noqa: BLE001
    raise SystemExit(f"SAM2 import failed from {sam2_libs}: {exc}") from exc
try:
    from utils.cotracker_wrist_grid import load_cotracker_predictor  # noqa: F401
except Exception as exc:  # noqa: BLE001
    raise SystemExit(f"CoTracker import failed from {openpi_root}: {exc}") from exc
PY

    IFS=',' read -r -a RAW_GPUS <<< "$GPU_LIST"
    GPUS_NORMALIZED=()
    declare -A SEEN_GPUS=()
    for gpu in "${RAW_GPUS[@]}"; do
        gpu="${gpu//[[:space:]]/}"
        if [[ -n "$gpu" ]]; then
            if [[ -n "${SEEN_GPUS[$gpu]:-}" ]]; then
                echo "ERROR: duplicate GPU in --gpus: $gpu" >&2
                exit 2
            fi
            SEEN_GPUS["$gpu"]=1
            GPUS_NORMALIZED+=("$gpu")
        fi
    done
    if [[ "${#GPUS_NORMALIZED[@]}" -eq 0 ]]; then
        echo "ERROR: --gpus must contain at least one GPU ID or UUID" >&2
        exit 1
    fi
    if ! command -v nvidia-smi >/dev/null 2>&1; then
        echo "ERROR: nvidia-smi is required for memory-aware scheduling" >&2
        exit 1
    fi
    for gpu in "${GPUS_NORMALIZED[@]}"; do
        gpu_total="$(
            nvidia-smi -i "$gpu" \
                --query-gpu=memory.total \
                --format=csv,noheader,nounits 2>/dev/null
        )" || {
            echo "ERROR: cannot query GPU: $gpu" >&2
            exit 1
        }
        gpu_total="${gpu_total//[[:space:]]/}"
        if [[ ! "$gpu_total" =~ ^[0-9]+$ ]]; then
            echo "ERROR: invalid memory.total returned for GPU $gpu: $gpu_total" >&2
            exit 1
        fi
        if ((GPU_MIN_FREE_MIB > gpu_total)); then
            echo "ERROR: --min-free-gpu-mib=$GPU_MIN_FREE_MIB exceeds GPU $gpu total ${gpu_total} MiB" >&2
            exit 2
        fi
    done

    IFS=',' read -r -a RAW_IDLE_GPUS <<< "$GPU_REQUIRE_IDLE"
    IDLE_GPUS_NORMALIZED=()
    for gpu in "${RAW_IDLE_GPUS[@]}"; do
        gpu="${gpu//[[:space:]]/}"
        if [[ -n "$gpu" ]]; then
            IDLE_GPUS_NORMALIZED+=("$gpu")
        fi
    done

    mkdir -p "$TRACK_CACHE"
    RUN_DIR="$DATASET_ROOT/meta/trex_track_force_prepare"
    mkdir -p "$RUN_DIR"
    if [[ "$START_EPISODE" == "auto" ]]; then
        if [[ "$FORCE_REBUILD" == "1" ]]; then
            RESUME_START=0
        else
            RESUME_START="$(
                "$PYTHON_BIN" - "$TRACK_CACHE" "$TOTAL_EPISODES" <<'PY'
import sys
import zipfile
from pathlib import Path

track_cache = Path(sys.argv[1])
total_episodes = int(sys.argv[2])
required_members = {
    "tracks.npy",
    "vis.npy",
    "track_layout_version.npy",
}

for episode_index in range(total_episodes):
    path = track_cache / f"episode_{episode_index:06d}.npz"
    if not path.is_file() or path.stat().st_size == 0:
        print(episode_index)
        break
    try:
        with zipfile.ZipFile(path) as archive:
            if not required_members.issubset(archive.namelist()):
                print(episode_index)
                break
    except (OSError, zipfile.BadZipFile):
        print(episode_index)
        break
else:
    print(total_episodes)
PY
            )"
        fi
    else
        RESUME_START="$START_EPISODE"
    fi
    if ((RESUME_START > TOTAL_EPISODES)); then
        echo "ERROR: start episode $RESUME_START exceeds total $TOTAL_EPISODES" >&2
        exit 2
    fi

    SESSION_ID="$(date -u +%Y%m%dT%H%M%SZ)_$$"
    SESSION_DIR="$RUN_DIR/run_$SESSION_ID"
    QUEUE_DIR="$SESSION_DIR/queue"
    mkdir -p "$QUEUE_DIR"

    echo "Preparing the complete T-Rex v2 dataset"
    echo "  episodes:    $TOTAL_EPISODES"
    echo "  output:      $DATASET_ROOT"
    echo "  track cache: $TRACK_CACHE"
    echo "  GPUs:        ${GPUS_NORMALIZED[*]}"
    echo "  force:       existing observation.tactile_force [60] -> [10,6]"
    echo "  tracks:      SAM2 + CoTracker -> [250,2] + visibility [250]"
    echo "  resume from: episode $RESUME_START"
    echo "  queue:       $EPISODES_PER_JOB episodes/task"
    echo "  memory gate: ${GPU_MIN_FREE_MIB} MiB free"
    if [[ "${#IDLE_GPUS_NORMALIZED[@]}" -gt 0 ]]; then
        echo "  idle gate:   ${IDLE_GPUS_NORMALIZED[*]}"
    fi

    for ((start = RESUME_START; start < TOTAL_EPISODES; start += EPISODES_PER_JOB)); do
        end=$((start + EPISODES_PER_JOB))
        if ((end > TOTAL_EPISODES)); then
            end="$TOTAL_EPISODES"
        fi
        printf -v task_name 'pending_%06d_%06d.task' "$start" "$end"
        : > "$QUEUE_DIR/$task_name"
    done

    queue_has_pending() {
        compgen -G "$QUEUE_DIR/pending_*.task" >/dev/null
    }

    gpu_requires_idle() {
        local candidate="$1"
        local idle_gpu
        for idle_gpu in "${IDLE_GPUS_NORMALIZED[@]:-}"; do
            if [[ "$candidate" == "$idle_gpu" ]]; then
                return 0
            fi
        done
        return 1
    }

    wait_for_gpu() {
        local gpu="$1"
        local min_free_mib="$2"
        local require_idle="$3"
        local stop_when_queue_empty="$4"
        local announced=0
        local free_mib process_output has_process
        while true; do
            if [[ "$stop_when_queue_empty" == "1" ]] && ! queue_has_pending; then
                return 2
            fi
            free_mib="$(
                nvidia-smi -i "$gpu" \
                    --query-gpu=memory.free \
                    --format=csv,noheader,nounits 2>/dev/null
            )" || {
                echo "ERROR: GPU $gpu memory query failed" >&2
                return 1
            }
            free_mib="${free_mib//[[:space:]]/}"
            if [[ ! "$free_mib" =~ ^[0-9]+$ ]]; then
                echo "ERROR: GPU $gpu returned invalid free memory: $free_mib" >&2
                return 1
            fi
            process_output="$(
                nvidia-smi -i "$gpu" \
                    --query-compute-apps=pid \
                    --format=csv,noheader,nounits 2>/dev/null
            )" || {
                echo "ERROR: GPU $gpu process query failed" >&2
                return 1
            }
            has_process=0
            if [[ -n "${process_output//[[:space:]]/}" ]]; then
                has_process=1
            fi
            if ((free_mib >= min_free_mib)) \
                && { [[ "$require_idle" == "0" ]] || [[ "$has_process" == "0" ]]; }; then
                if [[ "$announced" == "1" ]]; then
                    echo "GPU $gpu is ready (${free_mib} MiB free)"
                fi
                return 0
            fi
            if [[ "$announced" == "0" ]]; then
                if [[ "$require_idle" == "1" && "$has_process" == "1" ]]; then
                    echo "GPU $gpu is waiting for existing compute processes to exit"
                else
                    echo "GPU $gpu is waiting for ${min_free_mib} MiB free (now ${free_mib} MiB)"
                fi
                announced=1
            fi
            sleep "$GPU_POLL_SECONDS"
        done
    }

    claim_next_task() {
        local worker_id="$1"
        local candidate suffix claimed
        while true; do
            for candidate in "$QUEUE_DIR"/pending_*.task; do
                if [[ ! -e "$candidate" ]]; then
                    return 1
                fi
                suffix="${candidate##*/pending_}"
                claimed="$QUEUE_DIR/running_${worker_id}_${suffix}"
                if mv -- "$candidate" "$claimed" 2>/dev/null; then
                    printf '%s\n' "$claimed"
                    return 0
                fi
            done
        done
    }

    run_gpu_worker() {
        local worker_id="$1"
        local gpu="$2"
        local base_require_idle=0
        local wait_status task_path task_file task_suffix start end
        local task_name task_manifest task_log attempt task_min_free
        local task_require_idle command_status tee_status
        local -a command pipeline_status

        if gpu_requires_idle "$gpu"; then
            base_require_idle=1
        fi

        while queue_has_pending; do
            set +e
            wait_for_gpu "$gpu" "$GPU_MIN_FREE_MIB" "$base_require_idle" 1
            wait_status="$?"
            set -e
            if [[ "$wait_status" == "2" ]]; then
                return 0
            elif [[ "$wait_status" != "0" ]]; then
                return "$wait_status"
            fi

            task_path="$(claim_next_task "$worker_id")" || continue
            task_file="${task_path##*/}"
            task_suffix="${task_file#running_${worker_id}_}"
            task_suffix="${task_suffix%.task}"
            start="${task_suffix%%_*}"
            end="${task_suffix##*_}"
            task_name="task_${start}_${end}"
            task_manifest="$SESSION_DIR/${task_name}.json"
            task_log="$SESSION_DIR/${task_name}.log"

            command=(
                "$PYTHON_BIN" "$BUILDER"
                --dataset-root "$DATASET_ROOT"
                --track-cache "$TRACK_CACHE"
                --manifest-path "$task_manifest"
                --episode-range "$start" "$end"
                --extract-missing
                --no-update-metadata
                "$FK_FLAG"
                --calib-path "$CALIB_PATH"
                --openpi-root "$OPENPI_ROOT"
                --cotracker-checkpoint "$COTRACKER_CHECKPOINT"
                --cotracker-device cuda:0
                --sam2-model "$SAM2_MODEL"
                --sam2-device cuda:0
                --sam2-libs "$SAM2_LIBS"
                --no-save-viz
                --no-save-sam2-masks
            )
            if [[ "$FORCE_REBUILD" == "1" ]]; then
                command+=(--force)
            fi

            attempt=1
            task_min_free="$GPU_MIN_FREE_MIB"
            task_require_idle="$base_require_idle"
            while true; do
                if ((attempt > 1)); then
                    wait_for_gpu "$gpu" "$task_min_free" "$task_require_idle" 0
                fi
                echo "$task_name -> GPU $gpu (attempt $attempt; log: $task_log)" \
                    | tee -a "$task_log"
                set +e
                (
                    export CUDA_VISIBLE_DEVICES="$gpu"
                    export CUDA_MODULE_LOADING="${CUDA_MODULE_LOADING:-LAZY}"
                    export PYTHONUNBUFFERED=1
                    export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}"
                    "${command[@]}"
                ) 2>&1 | tee -a "$task_log"
                pipeline_status=("${PIPESTATUS[@]}")
                set -e
                command_status="${pipeline_status[0]}"
                tee_status="${pipeline_status[1]}"

                if [[ "$command_status" == "0" && "$tee_status" == "0" ]]; then
                    mv -- "$task_path" "$QUEUE_DIR/done_${task_suffix}.task"
                    break
                fi

                if [[ "$attempt" == "1" ]] \
                    && command -v rg >/dev/null 2>&1 \
                    && rg -q 'OutOfMemoryError|CUDA out of memory' "$task_log"; then
                    echo "WARNING: GPU $gpu hit OOM; retrying only after it is fully idle" \
                        | tee -a "$task_log" >&2
                    attempt=2
                    task_min_free="$GPU_RETRY_FREE_MIB"
                    task_require_idle=1
                    continue
                fi

                mv -- "$task_path" "$QUEUE_DIR/failed_${task_suffix}.task"
                if [[ "$command_status" == "0" ]]; then
                    command_status="$tee_status"
                fi
                echo "ERROR: $task_name failed on GPU $gpu" >&2
                return "$command_status"
            done
        done
    }

    PIDS=()
    WORKER_NAMES=()
    cleanup_workers() {
        local pid
        for pid in "${PIDS[@]:-}"; do
            if kill -0 "$pid" 2>/dev/null; then
                pkill -TERM -P "$pid" 2>/dev/null || true
                kill -TERM "$pid" 2>/dev/null || true
            fi
        done
    }
    trap 'cleanup_workers; exit 130' INT TERM HUP

    for worker in "${!GPUS_NORMALIZED[@]}"; do
        gpu="${GPUS_NORMALIZED[$worker]}"
        printf -v worker_name 'gpu_worker_%02d_gpu_%s' "$worker" "$gpu"
        echo "  $worker_name started"
        run_gpu_worker "$worker" "$gpu" &
        PIDS+=("$!")
        WORKER_NAMES+=("$worker_name")
    done

    failed=0
    for index in "${!PIDS[@]}"; do
        if ! wait "${PIDS[$index]}"; then
            echo "ERROR: ${WORKER_NAMES[$index]} failed; see $SESSION_DIR" >&2
            failed=1
        fi
    done
    trap - INT TERM HUP

    shopt -s nullglob
    unfinished_tasks=(
        "$QUEUE_DIR"/pending_*.task
        "$QUEUE_DIR"/running_*.task
        "$QUEUE_DIR"/failed_*.task
    )
    shopt -u nullglob
    if [[ "${#unfinished_tasks[@]}" -gt 0 ]]; then
        echo "Preprocessing stopped with ${#unfinished_tasks[@]} unfinished task(s)." >&2
        echo "Valid episodes and tracks are resumable; rerun the same command." >&2
        exit 1
    fi
    if [[ "$failed" == "1" ]]; then
        echo "WARNING: one or more GPU workers exited, but all queue tasks completed." >&2
    fi

    echo "Finalizing global metadata and canonical manifest..."
    "$PYTHON_BIN" "$BUILDER" \
        --dataset-root "$DATASET_ROOT" \
        --track-cache "$TRACK_CACHE" \
        --manifest-path "$CANONICAL_MANIFEST" \
        --all \
        --no-extract-missing \
        --update-metadata \
        "$FK_FLAG"

    echo "Running final all-episode validation..."
    "$PYTHON_BIN" "$BUILDER" \
        --dataset-root "$DATASET_ROOT" \
        --track-cache "$TRACK_CACHE" \
        --manifest-path "$CANONICAL_MANIFEST" \
        --all \
        --validate-only \
        "$FK_FLAG"

    echo "Prepared dataset: $DATASET_ROOT"
fi

if [[ "$RUN_TRAINING" == "1" ]]; then
    echo "Starting training with TREX_DATA_ROOT=$DATASET_ROOT"
    export TREX_DATA_ROOT="$DATASET_ROOT"
    export TRACK_CACHE="$TRACK_CACHE"
    export PATH="$(dirname "$PYTHON_BIN"):$PATH"
    exec bash "$TRAIN_SCRIPT" "${TRAIN_ARGS[@]}"
fi

cat <<EOF

Training command:
  TREX_DATA_ROOT="$DATASET_ROOT" \\
  TRACK_CACHE="$TRACK_CACHE" \\
  bash "$TRAIN_SCRIPT"
EOF