#!/usr/bin/env bash
# hf-fetch.sh — resumable, verified HuggingFace fetch (gated or public repos).
#
#   bash scripts/hf-fetch.sh <dest_dir> <manifest>
#
# Manifest: one line per file, three tab/space-separated fields:
#   <filename> <size_bytes> <sha256>
# A blank sha256 field means "check size only". Lines starting with '#' and
# blank lines are ignored. The first field is a path relative to <dest_dir>.
#
# Token: $HF_TOKEN, else ~/.cache/huggingface/token, else the llamaswap conf
# (grep hf_...). Passed out-of-band via a mode-600 curl -K file — a -H header
# is world-readable in /proc/<pid>/cmdline (the 2026-09-16 HF-token leak).
#
# Idempotent: a file whose size AND sha256 already match is skipped, so a
# re-run after an interruption only fetches what is missing or corrupt.
set -euo pipefail

dest="${1:?usage: hf-fetch.sh <dest_dir> <manifest>}"
manifest="${2:?usage: hf-fetch.sh <dest_dir> <manifest>}"
repo="${3:?usage: hf-fetch.sh <dest_dir> <manifest> <owner/repo>}"

mkdir -p "$dest"

# --- token resolution --------------------------------------------------------
token="${HF_TOKEN:-}"
if [[ -z "$token" && -r "$HOME/.cache/huggingface/token" ]]; then
    token="$(tr -d '[:space:]' < "$HOME/.cache/huggingface/token")"
fi
if [[ -z "$token" && -r "$HOME/.config/llamaswap.conf.yaml" ]]; then
    token="$(grep -oE 'hf_[a-zA-Z0-9]+' "$HOME/.config/llamaswap.conf.yaml" | head -1)"
fi
[[ -n "$token" ]] || { echo "hf-fetch: no HF token (set HF_TOKEN or ~/.cache/huggingface/token)" >&2; exit 3; }

auth_file="$(mktemp)"
chmod 600 "$auth_file"
printf 'header = "Authorization: Bearer %s"\n' "$token" > "$auth_file"
cleanup() { rm -f "$auth_file"; }
trap cleanup EXIT

base="https://huggingface.co/${repo}/resolve/main"

# --- per-file fetch + verify -------------------------------------------------
fetch_one() {
    local name="$1" size="$2" sha="$3"
    local url="$base/$name"
    local path="$dest/$name"
    mkdir -p "$(dirname "$path")"

    # idempotent skip
    if [[ -f "$path" ]]; then
        local cur_size sha_ok
        cur_size="$(wc -c < "$path")"
        sha_ok=true
        if [[ -n "$sha" && "$(sha256sum "$path" | cut -d' ' -f1)" != "$sha" ]]; then
            sha_ok=false
        fi
        if [[ "$cur_size" == "$size" && "$sha_ok" == "true" ]]; then
            echo "  ok   $name ($size bytes)"
            return 0
        fi
        echo "  refetch $name (have $cur_size, want $size)"
    fi

    # Resumable download; -f fails on HTTP error so a 401/403 is not written as a file.
    # -C - only when a non-empty partial exists: with a missing file it sends
    # Range: bytes=0-, which HF's xethub CDN rejects with 400 (xet-backed repos).
    local resume_args=()
    [[ -f "$path" && -s "$path" ]] && resume_args=(-C -)
    curl -fsSL "${resume_args[@]}" -K "$auth_file" -o "$path" "$url"

    # verify byte count
    local got
    got="$(wc -c < "$path")"
    if [[ "$got" != "$size" ]]; then
        echo "  FAIL $name: size $got != $size" >&2
        return 1
    fi
    # verify sha256
    if [[ -n "$sha" ]]; then
        local actual
        actual="$(sha256sum "$path" | cut -d' ' -f1)"
        if [[ "$actual" != "$sha" ]]; then
            echo "  FAIL $name: sha256 $actual != $sha" >&2
            return 1
        fi
    fi
    echo "  done $name ($size bytes)"
}

echo "hf-fetch: $repo -> $dest"
status=0
while IFS= read -r line || [[ -n "$line" ]]; do
    [[ -z "$line" || "$line" =~ ^[[:space:]]*# ]] && continue
    name="$(awk '{print $1}' <<<"$line")"
    size="$(awk '{print $2}' <<<"$line")"
    sha="$(awk '{print $3}' <<<"$line")"
    fetch_one "$name" "$size" "$sha" || status=1
done < "$manifest"

if [[ "$status" -ne 0 ]]; then
    echo "hf-fetch: FAILED (see above)" >&2
    exit 1
fi
echo "hf-fetch: all files verified"