#!/usr/bin/env bash
# Purpose: download one model file from the Hugging Face Hub with the hf CLI, verify its
#          SHA-256 against the checksum the Hub publishes for the file, and file it in a
#          shared model library at ~/models/<publisher>/<model>/ that every engine in this
#          course can point at
# Platform: all (Linux, macOS and WSL2; the hf CLI is identical on each)
# Minimum memory: 8 GB
# Assumes: hf, curl and python3 are on PATH; a sha256 tool (sha256sum or shasum) exists;
#          enough free disk for the file; HF_TOKEN is set or `hf auth login` has been run
#          if the repository is gated
#
# Usage: bash fetch-model.sh <repo-id> <filename> [labbook.md]
#   e.g. bash fetch-model.sh unsloth/Qwen3-8B-GGUF Qwen3-8B-Q4_K_M.gguf labbook.md
#
# Environment:
#   MODELS_DIR   where the library lives            (default: $HOME/models)
#   HF_REVISION  branch, tag or commit to download  (default: main)
#
# Re-running the script is safe and is the way to recover from an interruption: hf skips a
# file it has already recorded as complete and downloads one it has not, and this script
# re-verifies the checksum on every run. A file that fails the size or checksum test is
# moved aside together with hf's record of it, so the next run downloads it again.

set -euo pipefail

REPO="${1:-}"
FILE="${2:-}"
LABBOOK="${3:-}"
MODELS_DIR="${MODELS_DIR:-$HOME/models}"
HF_REVISION="${HF_REVISION:-main}"
API="https://huggingface.co/api/models"

die() { echo "fetch-model: $*" >&2; exit 1; }

[ -n "$REPO" ] && [ -n "$FILE" ] || die "usage: bash fetch-model.sh <repo-id> <filename> [labbook.md]"
case "$REPO" in
  */*) : ;;
  *) die "repo id must be <publisher>/<model>, got '$REPO'" ;;
esac

for tool in hf curl python3; do
  command -v "$tool" >/dev/null || die "$tool is not installed or not on PATH"
done

if command -v sha256sum >/dev/null; then
  sha256_of() { sha256sum "$1" | cut -d' ' -f1; }
elif command -v shasum >/dev/null; then
  sha256_of() { shasum -a 256 "$1" | cut -d' ' -f1; }
else
  die "no sha256 tool found (install coreutils for sha256sum, or use macOS shasum)"
fi

PUBLISHER="${REPO%%/*}"
MODEL="${REPO##*/}"
DEST="$MODELS_DIR/$PUBLISHER/$MODEL"

# --- 1. Ask the Hub what the file should be -----------------------------------------
# The repository tree endpoint reports, for every Git LFS file, an "lfs" object whose
# "oid" is the file's SHA-256 and whose "size" is its length in bytes.
echo "==> Looking up $FILE in $REPO@$HF_REVISION"
AUTH_ARGS=()
if [ -n "${HF_TOKEN:-}" ]; then
  AUTH_ARGS=(-H "Authorization: Bearer $HF_TOKEN")
fi

TREE_JSON="$(curl -fsSL ${AUTH_ARGS[@]+"${AUTH_ARGS[@]}"} "$API/$REPO/tree/$HF_REVISION" 2>/dev/null || true)"
EXPECTED=""
if [ -n "$TREE_JSON" ]; then
  EXPECTED="$(printf '%s' "$TREE_JSON" | python3 -c '
import json, sys
want = sys.argv[1]
try:
    entries = json.load(sys.stdin)
except json.JSONDecodeError:
    sys.exit(0)
for entry in entries:
    if entry.get("path") == want:
        lfs = entry.get("lfs") or {}
        if lfs.get("oid"):
            print(lfs["oid"], lfs.get("size", entry.get("size", 0)))
        break
' "$FILE")"
fi

# Fallback: the Git LFS pointer served at /raw/ carries the same SHA-256 and size.
if [ -z "$EXPECTED" ]; then
  echo "    tree listing gave nothing; falling back to the LFS pointer"
  POINTER="$(curl -fsSL ${AUTH_ARGS[@]+"${AUTH_ARGS[@]}"} "https://huggingface.co/$REPO/raw/$HF_REVISION/$FILE" || true)"
  OID="$(printf '%s\n' "$POINTER" | sed -n 's/^oid sha256:\([0-9a-f]\{64\}\)$/\1/p')"
  SIZE="$(printf '%s\n' "$POINTER" | sed -n 's/^size \([0-9]\{1,\}\)$/\1/p')"
  [ -n "$OID" ] && EXPECTED="$OID $SIZE"
fi

[ -n "$EXPECTED" ] || die "could not find a published SHA-256 for $FILE in $REPO (file names are case-sensitive: copy the name from the repository's file list; a private repository needs HF_TOKEN)"
EXPECTED_SHA="${EXPECTED%% *}"
EXPECTED_SIZE="${EXPECTED##* }"
echo "    expected sha256 $EXPECTED_SHA"
echo "    expected size   $EXPECTED_SIZE bytes"

# --- 2. Prepare the library ----------------------------------------------------------
mkdir -p "$DEST"
TARGET="$DEST/$FILE"
# hf keeps its record of each download here; removing it makes hf download the file again.
METADATA="$DEST/.cache/huggingface/download/$FILE.metadata"

# Refuse to start a download the disk cannot hold. df -Pk is POSIX, so the column is the
# same on Linux and macOS: available space in KiB.
HAVE_BYTES=0
[ -f "$TARGET" ] && HAVE_BYTES="$(python3 -c 'import os,sys; print(os.path.getsize(sys.argv[1]))' "$TARGET")"
NEED_KB=$(( (EXPECTED_SIZE - HAVE_BYTES) / 1024 ))
AVAIL_KB="$(df -Pk "$DEST" | awk 'NR == 2 { print $4 }')"
if [ "$NEED_KB" -gt 0 ] && [ "$AVAIL_KB" -lt "$NEED_KB" ]; then
  die "not enough free space on $DEST: need $(( NEED_KB / 1024 )) MiB more, $(( AVAIL_KB / 1024 )) MiB free. Free some space or set MODELS_DIR to another disk."
fi
SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
if [ ! -f "$MODELS_DIR/README.md" ]; then
  if [ -f "$SCRIPT_DIR/library-readme.md" ]; then
    cp "$SCRIPT_DIR/library-readme.md" "$MODELS_DIR/README.md"
  else
    printf '# Model library\n\nLayout: <publisher>/<model>/<file>. Each file has a .sha256 beside it.\n' \
      > "$MODELS_DIR/README.md"
  fi
  echo "==> Wrote $MODELS_DIR/README.md"
fi

# --- 3. Download ---------------------------------------------------------------------
echo "==> Downloading into $DEST (safe to interrupt and re-run)"
hf download "$REPO" "$FILE" --local-dir "$DEST" --revision "$HF_REVISION"

[ -f "$TARGET" ] || die "hf reported success but $TARGET does not exist"

# --- 4. Verify -----------------------------------------------------------------------
# hf decides whether a file is complete from its own metadata record, not from the bytes
# on disk, so a truncated or altered file with a record beside it would be kept forever.
# Anything that fails here is moved aside and its record removed, so a re-run downloads it.
set_aside() {
  mv "$TARGET" "$TARGET.corrupt"
  rm -f "$METADATA"
}

ACTUAL_SIZE="$(python3 -c 'import os,sys; print(os.path.getsize(sys.argv[1]))' "$TARGET")"
if [ "$ACTUAL_SIZE" != "$EXPECTED_SIZE" ]; then
  set_aside
  die "size mismatch: got $ACTUAL_SIZE bytes, expected $EXPECTED_SIZE. The file has been moved to $TARGET.corrupt; delete it and re-run to download it again."
fi

[ -f "$TARGET.sha256" ] && echo "    a previous run verified this file; verifying it again"
echo "==> Verifying SHA-256 (this reads the whole file; expect a minute or two on a large one)"
ACTUAL_SHA="$(sha256_of "$TARGET")"
if [ "$ACTUAL_SHA" != "$EXPECTED_SHA" ]; then
  set_aside
  die "checksum mismatch. Got $ACTUAL_SHA, expected $EXPECTED_SHA. The file has been moved to $TARGET.corrupt; delete it and re-run to download it again."
fi
printf '%s  %s\n' "$ACTUAL_SHA" "$FILE" > "$TARGET.sha256"
echo "    verified, checksum written to $TARGET.sha256"

# --- 5. Record -----------------------------------------------------------------------
if [ -n "$LABBOOK" ]; then
  # In a terminal `hf version` prints "✓ hf version" and "  version: x.y.z"; the agent
  # format is one "version=x.y.z" line. Ask for the agent format so the parse is stable.
  HF_VERSION="$(hf version --format agent 2>/dev/null | sed -n 's/^version=//p' | head -n 1)"
  python3 -c '
import json, sys
from datetime import date
print(json.dumps({
    "lab": "part-04/fetch-model",
    "date": date.today().isoformat(),
    "hf_version": sys.argv[7] or "unknown",
    "repo": sys.argv[1],
    "revision": sys.argv[2],
    "file": sys.argv[3],
    "path": sys.argv[4],
    "bytes": int(sys.argv[5]),
    "gb": round(int(sys.argv[5]) / 1e9, 2),
    "sha256": sys.argv[6],
    "verified": True,
}))' "$REPO" "$HF_REVISION" "$FILE" "$TARGET" "$ACTUAL_SIZE" "$ACTUAL_SHA" "$HF_VERSION" >> "$LABBOOK"
  echo "    recorded in $LABBOOK"
fi

echo "==> Done: $TARGET"