#!/usr/bin/env bash
# Full conversion driver for the new dataset mixture (run on eva01 from
# /mnt/storage/tinyvla_wds). Order minimizes peak disk: each phase downloads,
# converts to wds, pushes to AlexWortega/wds-*, validates, then deletes raw.
#
#   bash tinyvla_h100/scripts/convert_all.sh 2>&1 | tee logs/convert_all.log
set -uo pipefail
cd /mnt/storage/tinyvla_wds
export HF_HOME=/mnt/storage/tinyvla_wds/.hf_home
source venv/bin/activate
S=tinyvla_h100/scripts
mkdir -p data/wds raw logs
FAILED=""

phase() { echo; echo "=== [$(date +%H:%M:%S)] $1 ==="; df -h /mnt | tail -1; }
push() { # push <local_dir> <repo>
  python $S/hf_push_wds.py --src "$1" --repo "$2" || { FAILED="$FAILED $2(push)"; return 1; }
}

# ---------------------------------------------------------------- tartandrive
if [ ! -f data/wds/tartandrive/manifest.json ]; then
phase "tartandrive (vamos, 15GB, evict)"
python $S/vamos_to_wds.py --repo-id mateoguaman/tartandrive \
  --out data/wds/tartandrive --name tartandrive --embodiment-id 17 \
  --workers 24 --evict --instruction-fallback "drive along the off-road trail" \
  && python $S/validate_wds.py data/wds/tartandrive --batches 1 \
  && push data/wds/tartandrive AlexWortega/wds-tartandrive \
  || FAILED="$FAILED tartandrive"
fi

# ---------------------------------------------------------------------- scand
if [ ! -f data/wds/scand/manifest.json ]; then
phase "scand (vamos, subset 30 files, evict)"
python $S/vamos_to_wds.py --repo-id mateoguaman/scand \
  --out data/wds/scand --name scand --embodiment-id 18 \
  --max-files 30 --workers 24 --evict \
  --instruction-fallback "navigate the campus to the goal" \
  && python $S/validate_wds.py data/wds/scand --batches 1 \
  && push data/wds/scand AlexWortega/wds-scand \
  || FAILED="$FAILED scand"
fi

# --------------------------------------------------------------------- sacson
if [ ! -f data/wds/sacson/manifest.json ]; then
phase "sacson (bags, 42GB)"
python - <<'EOF'
from huggingface_hub import snapshot_download
snapshot_download("franciszzj/SACSoN", repo_type="dataset",
                  local_dir="raw/sacson", allow_patterns=["*.bag"])
EOF
python $S/gnm_to_wds.py --kind bags --root raw/sacson \
  --out data/wds/sacson --name sacson --embodiment-id 16 --hz 4 \
  --instruction "navigate around pedestrians to the goal" \
  && python $S/validate_wds.py data/wds/sacson --batches 1 \
  && push data/wds/sacson AlexWortega/wds-sacson \
  || FAILED="$FAILED sacson"
rm -rf raw/sacson .hf_home/hub/datasets--franciszzj--SACSoN
fi

# ---------------------------------------------------------------------- recon
if [ ! -f data/wds/recon/manifest.json ]; then
phase "recon (stream-extract 53GB tar.gz -> h5)"
TOK=$(cat ~/.cache/huggingface/token)
python - <<'EOF' > /tmp/recon_urls.txt
from huggingface_hub import HfApi
fs = sorted(f for f in HfApi().list_repo_files("franciszzj/RECON", repo_type="dataset")
            if f.startswith("recon_dataset_"))
for f in fs:
    print(f"https://huggingface.co/datasets/franciszzj/RECON/resolve/main/{f}")
EOF
mkdir -p raw/recon
while read -r u; do curl -sL -H "Authorization: Bearer $TOK" "$u"; done < /tmp/recon_urls.txt \
  | tar xzf - -C raw/recon
python $S/gnm_to_wds.py --kind recon_h5 --root raw/recon/recon_release \
  --out data/wds/recon --name recon --embodiment-id 16 \
  --instruction "explore the outdoor terrain to reach the goal" \
  && python $S/validate_wds.py data/wds/recon --batches 1 \
  && push data/wds/recon AlexWortega/wds-recon \
  || FAILED="$FAILED recon"
rm -rf raw/recon
fi

# ---------------------------------------------------------------- go_stanford
if [ ! -f data/wds/go_stanford/manifest.json ]; then
phase "go_stanford (already extracted, 3.5GB)"
python $S/gnm_to_wds.py --kind gnm_dirs --root raw/go_stanford \
  --out data/wds/go_stanford --name go_stanford --embodiment-id 16 \
  --instruction "navigate the indoor corridors to the goal" \
  && python $S/validate_wds.py data/wds/go_stanford --batches 1 \
  && push data/wds/go_stanford AlexWortega/wds-go-stanford \
  || FAILED="$FAILED go_stanford"
fi

# -------------------------------------------------------------------- unitree
phase "unitree (14 v3.0 sets)"
mkdir -p data/wds/unitree
UNITREE_G1="G1_Dex1_StackRygBlock_Dataset_Sim G1_Dex1_DiverseManip_DualArm_256x256 \
G1_Dex1_MountCamera_Dataset G1_Dex3_ToastedBread_Dataset G1_Dex1_DiverseManip_SingleArm_256x256 \
G1_Brainco_PickDoll_Dataset G1_Dex3_PickDoll_Dataset G1_Dex3_BlockStacking_Dataset \
G1_Dex3_GraspSquare_Dataset G1_Dex3_CameraPackaging_Dataset G1_Brainco_GraspOreo_Dataset"
UNITREE_Z1="Z1_StackBox_Dataset Z1_Dual_Dex1_PourCoffee_Dataset Z1_Dual_Dex1_FoldClothes_Dataset"
for name in $UNITREE_G1 $UNITREE_Z1; do
  [ -f data/wds/unitree/$name/manifest.json ] && continue
  case " $UNITREE_Z1 " in *" $name "*) EMB=19;; *) EMB=15;; esac
  echo "--- unitree/$name (embodiment $EMB)"
  python $S/build_shards.py --repo-id unitreerobotics/$name \
      --out data/shards_tmp/$name --hz 10 --cache-dir .hf_home/tmpcache \
    && python $S/pack_wds.py --src data/shards_tmp/$name \
      --out data/wds/unitree/$name --chunk 50 --maxcount 3000 --embodiment-id $EMB \
    && python $S/validate_wds.py data/wds/unitree/$name --batches 0 \
    || FAILED="$FAILED unitree/$name"
  rm -rf data/shards_tmp/$name .hf_home/tmpcache
done
push data/wds/unitree AlexWortega/wds-unitree || true

echo
echo "=== DONE. FAILED:${FAILED:-none} ==="
df -h /mnt | tail -1