#!/bin/bash
#
# Prepare resized TACO dataset for HuggingFace upload.
#
# 1. Zips each modality directory
# 2. Splits zips > 50GB into 49GB chunks
# 3. Generates MANIFEST.txt
#
# Usage: ./prepare_resized_upload.sh
#
set -euo pipefail

DST="${1:-/valhalla/projects/ehpc-dev-2025d08-044/taco_dataset_resized}"
cd "$DST"

THRESHOLD=$((50 * 1024 * 1024 * 1024))  # 50 GB
CHUNK="49G"
BACKUP_DIR="${DST}/_originals_backup"

# Directories to zip (each becomes a .zip at root level)
MODALITIES=(
    Marker_Removed_Allocentric_RGB_Videos
    2D_Segmentation
    Allocentric_Camera_Parameters
    Egocentric_RGB_Videos
    Egocentric_Depth_Videos
    Egocentric_Camera_Parameters
    Hand_Poses
    Object_Poses
    object_models_released
    mano_v1_2
)

echo "=== Step 1: Zip modality directories ==="
for mod in "${MODALITIES[@]}"; do
    if [ ! -d "$mod" ]; then
        echo "  SKIP $mod (not found)"
        continue
    fi

    zip_name="${mod}.zip"
    if [ -f "$zip_name" ]; then
        echo "  SKIP $mod (zip already exists)"
        continue
    fi

    echo "  Zipping $mod ..."
    zip -r -0 "$zip_name" "$mod"  # -0 = store only (no compression, faster)
    size=$(du -h "$zip_name" | cut -f1)
    echo "    Created $zip_name ($size)"
done
echo ""

echo "=== Step 2: Split zips > 50GB ==="
LARGE_FILES=()
while IFS= read -r -d '' f; do
    size=$(stat -c%s "$f")
    if [ "$size" -gt "$THRESHOLD" ]; then
        LARGE_FILES+=("$f")
    fi
done < <(find . -maxdepth 1 -type f -name '*.zip' -print0)

if [ ${#LARGE_FILES[@]} -eq 0 ]; then
    echo "  No zips exceed 50GB. Nothing to split."
else
    echo "  Found ${#LARGE_FILES[@]} files to split:"
    for f in "${LARGE_FILES[@]}"; do
        echo "    $f ($(du -h "$f" | cut -f1))"
    done
    echo ""

    mkdir -p "$BACKUP_DIR"
    for f in "${LARGE_FILES[@]}"; do
        echo "  [SPLIT] $f"
        split -b "$CHUNK" -d --additional-suffix=.part "$f" "${f}."
        echo "    Parts created:"
        ls -lh "${f}".*.part
        mv "$f" "$BACKUP_DIR/"
        echo "    Original moved to $BACKUP_DIR/"
        echo ""
    done
fi
echo ""

echo "=== Step 3: Clean up extracted directories ==="
echo "  (Keeping directories intact — remove manually after upload if desired)"
echo ""

echo "=== Step 4: Generate MANIFEST.txt ==="
{
    # List zips, parts, and standalone files
    find . -maxdepth 1 -type f \( -name '*.zip' -o -name '*.part' -o -name '*.csv' -o -name '*.md' -o -name '*.sh' -o -name '*.txt' \) | sort | while read -r f; do
        size=$(stat -c%s "$f")
        echo "$size  $f"
    done
} > MANIFEST.txt
echo "  Wrote MANIFEST.txt ($(wc -l < MANIFEST.txt) entries)"
cat MANIFEST.txt
echo ""

echo "=== Done ==="
echo "To upload:"
echo "  huggingface-cli upload mzhobro/taco_dataset_resized $DST --repo-type dataset"