#!/usr/bin/env bash
set -euo pipefail

usage() {
  echo "usage: cache-warm pnpm LOCKFILE... | venv-key INPUT... | venv-restore VENV INPUT... | venv-save VENV INPUT..." >&2
  exit 64
}

hash_inputs() {
  (( $# > 0 )) || usage
  local input
  for input in "$@"; do
    [[ -f "$input" ]] || { echo "cache-warm: missing input: $input" >&2; return 66; }
  done
  # Include stable logical names and bytes; input order and parent path do not
  # affect the key. Sorting the complete records also handles repeated basenames.
  for input in "$@"; do
    printf '%s:%s\0' "${input##*/}" "$(sha256sum "$input" | cut -d ' ' -f 1)"
  done | sort -z | sha256sum | cut -d ' ' -f 1
}

archive_for() {
  printf '%s/venv-%s.tar.zst\n' "${VENV_ARCHIVE_DIR:-/cache/venv-archives}" "$(hash_inputs "$@")"
}

case "${1:-}" in
  pnpm)
    shift; (( $# > 0 )) || usage
    export FT_FROM_HOOK=1 FT_HINTS=1
    for lockfile in "$@"; do
      [[ -f "$lockfile" ]] || { echo "cache-warm: missing lockfile: $lockfile" >&2; exit 66; }
      pnpm fetch --frozen-lockfile --lockfile-dir "$(dirname "$lockfile")" \
        --store-dir "${PNPM_STORE_DIR:-/cache/pnpm/store}"
    done
    ;;
  venv-key)
    shift; hash_inputs "$@"
    ;;
  venv-restore)
    shift; (( $# >= 2 )) || usage
    venv=$1; shift; archive=$(archive_for "$@")
    [[ -f "$archive" ]] || { echo "cache-warm: miss ${archive##*/}"; exit 3; }
    [[ ! -e "$venv" ]] || { echo "cache-warm: destination exists: $venv" >&2; exit 65; }
    mkdir -p "$venv"
    tar --zstd -xf "$archive" -C "$venv"
    echo "cache-warm: restored ${archive##*/}"
    ;;
  venv-save)
    shift; (( $# >= 2 )) || usage
    venv=$1; shift; [[ -d "$venv" ]] || { echo "cache-warm: missing venv: $venv" >&2; exit 66; }
    archive=$(archive_for "$@"); mkdir -p "$(dirname "$archive")"
    if [[ -f "$archive" ]]; then echo "cache-warm: already saved ${archive##*/}"; exit 0; fi
    tmp="${archive}.tmp.$$"
    trap 'rm -f "$tmp"' EXIT
    tar --zstd -cf "$tmp" -C "$venv" .
    # A hard link publishes the complete archive atomically; first writer wins.
    if ln "$tmp" "$archive" 2>/dev/null; then rm -f "$tmp"; else rm -f "$tmp"; fi
    trap - EXIT
    echo "cache-warm: saved ${archive##*/}"
    ;;
  *) usage ;;
esac