#!/usr/bin/env bash
# Purpose: take a restorable backup of everything about your model estate that is not the
#          weights - configurations, adapters, evaluation sets, the lab notebook, the
#          dashboard, and a manifest recording every model file's hash - verify the
#          archive can be read back, and record what was taken
# Platform: all (Linux, macOS and WSL2)
# Minimum memory: 8 GB, which is the service being backed up; this script needs almost none
# Assumes: tar on PATH, python3 for the manifest step, and enough free space in the
#          destination. Model weights are excluded on purpose: they are large, they are
#          re-downloadable, and the manifest records the hash of each one so you can prove
#          that what you download next year is what you had this year. Pass
#          --include-weights when you have somewhere to put them and a reason.
#
# Usage: bash backup-estate.sh --gateway ~/gateway --models ~/models --dest ~/backups
#        bash backup-estate.sh --gateway ~/gateway --models ~/models --dest ~/backups \
#             --extra ~/adapters --extra ~/evals --labbook labbook.md
set -euo pipefail

GATEWAY_DIR=""
MODELS_DIR=""
DEST="${HOME}/backups"
LABBOOK=""
INCLUDE_WEIGHTS=0
EXTRAS=()

HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
ENV_FILE="${ENV_FILE:-${HERE}/.env}"
if [ -f "$ENV_FILE" ]; then
    set -a
    # shellcheck source=/dev/null
    . "$ENV_FILE"
    set +a
    DEST="${BACKUP_DIR:-$DEST}"
fi

while [ $# -gt 0 ]; do
    case "$1" in
        --gateway)         GATEWAY_DIR="${2:?--gateway needs a directory}"; shift 2 ;;
        --models)          MODELS_DIR="${2:?--models needs a directory}"; shift 2 ;;
        --dest)            DEST="${2:?--dest needs a directory}"; shift 2 ;;
        --extra)           EXTRAS+=("${2:?--extra needs a directory}"); shift 2 ;;
        --labbook)         LABBOOK="${2:?--labbook needs a file}"; shift 2 ;;
        --include-weights) INCLUDE_WEIGHTS=1; shift ;;
        -h|--help)         sed -n '2,22p' "$0"; exit 0 ;;
        *) echo "unknown argument: $1" >&2; exit 2 ;;
    esac
done

command -v tar >/dev/null 2>&1 || { echo "tar is not on PATH." >&2; exit 1; }
[ -n "$GATEWAY_DIR" ] || { echo "Give --gateway <directory>: the Part 9 gateway directory." >&2; exit 1; }
[ -d "$GATEWAY_DIR" ] || { echo "No such directory: ${GATEWAY_DIR}" >&2; exit 1; }

mkdir -p "$DEST"
STAMP="$(date -u +%Y%m%dT%H%M%SZ)"
WORK="$(mktemp -d)"
trap 'rm -rf "$WORK"' EXIT
ARCHIVE="${DEST}/estate-${STAMP}.tar.gz"

echo "==> Writing the manifest"
MANIFEST="${WORK}/estate-manifest.json"
if [ -f "${HERE}/estate-manifest.py" ] && command -v python3 >/dev/null 2>&1; then
    MANIFEST_ARGS=(--gateway "$GATEWAY_DIR" --output "$MANIFEST")
    [ -n "$MODELS_DIR" ] && MANIFEST_ARGS+=(--models "$MODELS_DIR")
    python3 "${HERE}/estate-manifest.py" "${MANIFEST_ARGS[@]}"
else
    echo "    estate-manifest.py is not beside this script; skipping the manifest." >&2
    printf '{"note":"no manifest was written"}\n' > "$MANIFEST"
fi

echo "==> Collecting"
STAGE="${WORK}/estate"
mkdir -p "$STAGE"
cp "$MANIFEST" "${STAGE}/estate-manifest.json"

# The gateway directory, minus the one file in it that must never leave the machine.
mkdir -p "${STAGE}/gateway"
tar -cf - -C "$(dirname "$GATEWAY_DIR")" \
    --exclude='.env' --exclude='*.key' --exclude='*.pem' \
    "$(basename "$GATEWAY_DIR")" | tar -xf - -C "${STAGE}/gateway"

for extra in ${EXTRAS+"${EXTRAS[@]}"}; do
    if [ -d "$extra" ]; then
        echo "    including ${extra}"
        mkdir -p "${STAGE}/extra"
        tar -cf - -C "$(dirname "$extra")" "$(basename "$extra")" | tar -xf - -C "${STAGE}/extra"
    else
        echo "    skipping ${extra}: not a directory" >&2
    fi
done

if [ -n "$LABBOOK" ] && [ -f "$LABBOOK" ]; then
    cp "$LABBOOK" "${STAGE}/labbook.md"
fi

if [ "$INCLUDE_WEIGHTS" -eq 1 ]; then
    [ -n "$MODELS_DIR" ] || { echo "--include-weights needs --models." >&2; exit 1; }
    echo "    including the model library, which will take a while and a lot of space"
    mkdir -p "${STAGE}/models"
    tar -cf - -C "$(dirname "$MODELS_DIR")" "$(basename "$MODELS_DIR")" \
        | tar -xf - -C "${STAGE}/models"
fi

echo "==> Archiving"
tar -czf "$ARCHIVE" -C "$WORK" estate

echo "==> Verifying the archive can be read back"
ENTRIES="$(tar -tzf "$ARCHIVE" | wc -l | tr -d ' ')"
BYTES="$(wc -c < "$ARCHIVE" | tr -d ' ')"
if [ "$ENTRIES" -lt 2 ]; then
    echo "The archive contains almost nothing; something went wrong." >&2
    exit 1
fi
if ! tar -tzf "$ARCHIVE" | grep -q 'estate/estate-manifest.json'; then
    echo "The manifest is not in the archive; something went wrong." >&2
    exit 1
fi

echo "    ${ARCHIVE}"
echo "    ${ENTRIES} entries, ${BYTES} bytes"
echo "    .env, private keys and certificates were excluded on purpose. Keep those"
echo "    somewhere else, deliberately, and know where."

if [ -n "$LABBOOK" ]; then
    printf '{"lab":"part-23/backup","archive":"%s","entries":%s,"bytes":%s,"weights_included":%s,"taken":"%s"}\n' \
        "$(basename "$ARCHIVE")" "$ENTRIES" "$BYTES" "$INCLUDE_WEIGHTS" "$STAMP" >> "$LABBOOK"
    echo "    recorded in ${LABBOOK}"
fi

echo "==> Done. Now restore it somewhere harmless and check that it is what you think."