#!/bin/bash
# Download the TLT12 dataset.
#
#   ./download.sh                        # everything, into ./TLT12
#   ./download.sh --split val            # just the val split (try before you commit)
#   ./download.sh --dest /data/TLT12     # somewhere else
#   ./download.sh --jobs 4               # 4 parallel transfers
#   ./download.sh --list                 # show what would be downloaded, and how big
#
# Safe to re-run: files already present with the right checksum are skipped
# without touching the network, and a partial file is resumed rather than
# refetched.  So if it is interrupted, just run the same command again.
#
# The file list comes from SHA256SUMS, which is fetched first and also used to
# verify every file afterwards.
set -euo pipefail

# TODO: set this to the published location before release.
BASE_URL="${TLT12_BASE_URL:-https://zenkelab.org/datasets/TLT12}"

DEST="./TLT12"
SPLIT=""
JOBS=1
VERIFY=1
LIST_ONLY=0

usage() {
    sed -n '2,16p' "$0" | sed 's/^# \?//'
    cat <<'EOF'

Options:
  --base-url URL   Where the release is hosted (or set $TLT12_BASE_URL)
  --dest DIR       Download destination (default: ./TLT12)
  --split NAME     Only fetch one split: train, val, test, or ood-test
  --jobs N         Parallel transfers (default: 1)
  --no-verify      Skip the final checksum pass
  --list           Print the file list and total size, download nothing
  -h, --help       This message
EOF
}

while [[ $# -gt 0 ]]; do
    case "$1" in
        --base-url) BASE_URL="$2"; shift 2 ;;
        --dest)     DEST="$2";     shift 2 ;;
        --split)    SPLIT="$2";    shift 2 ;;
        --jobs)     JOBS="$2";     shift 2 ;;
        --no-verify) VERIFY=0;     shift ;;
        --list)     LIST_ONLY=1;   shift ;;
        -h|--help)  usage; exit 0 ;;
        *) echo "Unknown option: $1" >&2; usage >&2; exit 2 ;;
    esac
done

command -v curl     >/dev/null || { echo "ERROR: curl is required" >&2; exit 1; }
command -v sha256sum >/dev/null || { echo "ERROR: sha256sum is required" >&2; exit 1; }

BASE_URL="${BASE_URL%/}"
DEST="${DEST%/}"

# ── fetch the manifest ───────────────────────────────────────────────────────
mkdir -p "$DEST"
echo "Fetching manifest from $BASE_URL/SHA256SUMS"
if ! curl -fsSL --retry 3 -o "$DEST/SHA256SUMS" "$BASE_URL/SHA256SUMS"; then
    echo "ERROR: could not fetch $BASE_URL/SHA256SUMS" >&2
    echo "       Check --base-url (currently: $BASE_URL)" >&2
    exit 1
fi

# ── decide what to fetch ─────────────────────────────────────────────────────
# Manifest lines are: <64 hex chars><2 spaces><relative path>
# Everything outside shards/ is small (docs, loader, metadata) and is always
# fetched, so even a --split download is a usable release.
SUBSET="$DEST/.manifest_subset"
if [[ -n "$SPLIT" ]]; then
    if ! awk '{ print substr($0, 67) }' "$DEST/SHA256SUMS" \
         | grep -q "^shards/$SPLIT/"; then
        echo "ERROR: no split named '$SPLIT' in the manifest. Available:" >&2
        awk '{ print substr($0, 67) }' "$DEST/SHA256SUMS" \
            | grep '^shards/' | cut -d/ -f2 | sort -u | sed 's/^/       /' >&2
        exit 1
    fi
    awk -v want="^shards/$SPLIT/" \
        '{ p = substr($0, 67); if (p !~ /^shards\// || p ~ want) print }' \
        "$DEST/SHA256SUMS" > "$SUBSET"
else
    cp "$DEST/SHA256SUMS" "$SUBSET"
fi

N_FILES=$(grep -c . "$SUBSET" || true)
echo "$N_FILES files in manifest${SPLIT:+ (split: $SPLIT)}"

if [[ "$LIST_ONLY" -eq 1 ]]; then
    awk '{ print "  " substr($0, 67) }' "$SUBSET"
    echo
    echo "Querying total size ..."
    awk '{ print substr($0, 67) }' "$SUBSET" | while read -r f; do
        curl -fsSLI "$BASE_URL/$f" \
            | awk 'tolower($1) == "content-length:" { print $2 }' | tail -1
    done | awk '{ t += $1 } END { printf "  %.2f GB across %d files\n", t / 1e9, NR }'
    rm -f "$SUBSET"
    exit 0
fi

# ── download ─────────────────────────────────────────────────────────────────
# Three cases per file:
#   present and correct -> skip entirely, no request made
#   present but wrong   -> partial download; resume, refetch whole if the
#                          server has no range support
#   absent              -> plain fetch
fetch_one() {
    local hash="$1" rel="$2" base="$3" dest="$4"
    local out="$dest/$rel"

    if [[ -f "$out" ]] && [[ "$(sha256sum "$out" | cut -d' ' -f1)" == "$hash" ]]; then
        echo "  have  $rel"
        return 0
    fi

    mkdir -p "$(dirname "$out")"

    if [[ -f "$out" ]]; then
        if curl -fsSL --retry 5 --retry-delay 2 -C - -o "$out" "$base/$rel"; then
            echo "  resumed  $rel"
            return 0
        fi
        echo "  (cannot resume $rel — refetching in full)" >&2
        rm -f "$out"
    fi

    if curl -fsSL --retry 5 --retry-delay 2 -o "$out" "$base/$rel"; then
        echo "  ok    $rel"
        return 0
    fi
    echo "FAILED: $rel" >&2
    return 1
}
export -f fetch_one
export BASE_URL DEST

echo "Downloading to $DEST"
if [[ "$JOBS" -gt 1 ]]; then
    xargs -P "$JOBS" -n 2 \
        bash -c 'fetch_one "$0" "$1" "$BASE_URL" "$DEST"' < "$SUBSET"
else
    while read -r hash path; do
        [[ -z "${path:-}" ]] && continue
        fetch_one "$hash" "$path" "$BASE_URL" "$DEST"
    done < "$SUBSET"
fi

# ── verify ───────────────────────────────────────────────────────────────────
if [[ "$VERIFY" -eq 1 ]]; then
    echo
    echo "Verifying checksums ..."
    if (cd "$DEST" && sha256sum -c --quiet "$(basename "$SUBSET")"); then
        echo "All checksums OK."
    else
        echo "ERROR: checksum failure above. Re-run this script — it refetches" >&2
        echo "       only the bad files and leaves the good ones alone." >&2
        rm -f "$SUBSET"
        exit 1
    fi
fi
rm -f "$SUBSET"

cat <<EOF

Done. The dataset is in $DEST

Next:
    cd $DEST
    pip install -r requirements.txt
    python verify_release.py
    python example_load.py --shards shards/${SPLIT:-train}
EOF