#!/usr/bin/env bash
# Purpose: make ~/models/hf the Hugging Face home (HF_HOME) on this machine, idempotently:
#          create it inside the model library Part 4 created, add HF_HOME to your shell
#          profile once, copy an existing Hugging Face login across, and optionally move
#          an existing hub cache into it
# Platform: all (Linux, WSL2 and macOS; bash or zsh as the login shell)
# Minimum memory: 8 GB
# Assumes: Part 4's library at ~/models (created here if it is missing); SHELL names bash
#          or zsh; nothing is ever deleted, and nothing already in place is overwritten
#
# Usage: bash setup-model-storage.sh [--check] [--move-cache]
#        --check       report what would change, and change nothing
#        --move-cache  also move an existing ~/.cache/huggingface/hub to ~/models/hf/hub,
#                      only if ~/models/hf/hub does not exist yet
#        --help        print this header
#
# Running it twice is safe: the second run reports "ok" for everything the first created.
set -euo pipefail

MODELS="${MODELS:-$HOME/models}"
NEW_HOME="$MODELS/hf"
OLD_HOME="${XDG_CACHE_HOME:-$HOME/.cache}/huggingface"
CHECK="no"
MOVE_CACHE="no"

for arg in "$@"; do
  case "$arg" in
    --check) CHECK="yes" ;;
    --move-cache) MOVE_CACHE="yes" ;;
    -h|--help) sed -n '2,17p' "$0"; exit 0 ;;
    *) echo "setup-model-storage: unknown argument '$arg' (try --help)" >&2; exit 2 ;;
  esac
done

CHANGED=0
NOTES=0
ok()      { printf '   ok       %-20s %s\n' "$1" "$2"; }
changed() { CHANGED=$((CHANGED + 1)); printf '   %-8s %-20s %s\n' "$ACTION" "$1" "$2"; }
note()    { NOTES=$((NOTES + 1)); printf '   NOTE     %-20s %s\n' "$1" "$2"; }
if [ "$CHECK" = "yes" ]; then ACTION="WOULD"; else ACTION="DONE"; fi

# The line written to the profile. With the default library it stays portable ($HOME).
if [ "$MODELS" = "$HOME/models" ]; then
  # shellcheck disable=SC2016  # the literal $HOME is intended: the profile expands it
  PROFILE_LINE='export HF_HOME="$HOME/models/hf"'
else
  PROFILE_LINE="export HF_HOME=\"$NEW_HOME\""
fi

case "$(basename "${SHELL:-}")" in
  zsh) PROFILE="$HOME/.zshrc" ;;
  bash)
    if [ "$(uname -s)" = "Darwin" ]; then PROFILE="$HOME/.bash_profile"; else PROFILE="$HOME/.bashrc"; fi ;;
  *)
    echo "setup-model-storage: SHELL is '${SHELL:-unset}', not bash or zsh." >&2
    echo "Add this line to your shell's startup file yourself, then rerun with --check:" >&2
    echo "  $PROFILE_LINE" >&2
    exit 1 ;;
esac

printf 'setup-model-storage.sh, mode: %s\n' "$([ "$CHECK" = "yes" ] && echo check || echo apply)"
printf '   library  %s\n   HF_HOME  %s\n   profile  %s\n\n' "$MODELS" "$NEW_HOME" "$PROFILE"

# 1. The directory, inside Part 4's library
if [ -d "$NEW_HOME" ]; then
  ok "directory" "$NEW_HOME exists"
else
  [ "$CHECK" = "yes" ] || mkdir -p "$NEW_HOME"
  changed "directory" "create $NEW_HOME"
fi
if [ -f "$MODELS/README.md" ]; then
  ok "library README" "$MODELS/README.md (Part 4)"
else
  note "library README" "missing: Part 4's lab, task 7, puts it there; the layout still works"
fi

# 2. The profile line, added once
if [ -f "$PROFILE" ] && grep -Eq '^[[:space:]]*export[[:space:]]+HF_HOME=' "$PROFILE"; then
  existing="$(grep -E '^[[:space:]]*export[[:space:]]+HF_HOME=' "$PROFILE" | tail -n 1 | sed 's/^[[:space:]]*//')"
  if [ "$existing" = "$PROFILE_LINE" ]; then
    ok "profile" "$PROFILE already sets HF_HOME"
  else
    note "profile" "$PROFILE already has: $existing (left unchanged; edit it by hand if it is wrong)"
  fi
else
  if [ "$CHECK" = "no" ]; then
    printf '\n# Hugging Face home inside the model library (Local LLM course, Part 5)\n%s\n' \
      "$PROFILE_LINE" >> "$PROFILE"
  fi
  changed "profile" "append to $PROFILE: $PROFILE_LINE"
fi

# 3. An existing login: hf reads the token from $HF_HOME/token, so copy it across
for f in token stored_tokens; do
  if [ -e "$NEW_HOME/$f" ]; then
    ok "login: $f" "present in $NEW_HOME"
  elif [ -f "$OLD_HOME/$f" ]; then
    [ "$CHECK" = "yes" ] || install -m 600 "$OLD_HOME/$f" "$NEW_HOME/$f"
    changed "login: $f" "copy $OLD_HOME/$f (mode 600)"
  else
    ok "login: $f" "none stored under $OLD_HOME, nothing to copy"
  fi
done

# 4. An existing hub cache: reported, and moved only when asked
if [ -d "$OLD_HOME/hub" ]; then
  size="$(du -sh "$OLD_HOME/hub" 2>/dev/null | cut -f1)"
  if [ -e "$NEW_HOME/hub" ]; then
    note "old hub cache" "$OLD_HOME/hub ($size) and $NEW_HOME/hub both exist; not merging them"
  elif [ "$MOVE_CACHE" = "yes" ]; then
    [ "$CHECK" = "yes" ] || mv "$OLD_HOME/hub" "$NEW_HOME/hub"
    changed "old hub cache" "move $OLD_HOME/hub ($size) to $NEW_HOME/hub"
  else
    note "old hub cache" "$OLD_HOME/hub holds $size; rerun with --move-cache to move it"
  fi
else
  ok "old hub cache" "none at $OLD_HOME/hub"
fi

# 5. This shell
if [ "${HF_HOME:-}" = "$NEW_HOME" ]; then
  ok "this shell" "HF_HOME=$HF_HOME"
else
  note "this shell" "HF_HOME is '${HF_HOME:-unset}'; open a new terminal before using hf"
fi

printf '\nsummary: %d %s, %d NOTE\n' "$CHANGED" "$([ "$CHECK" = "yes" ] && echo "would change" || echo changed)" "$NOTES"