#!/usr/bin/env bash
# ~/docker/ai-image-gen/archive.sh
#
# Run this once after `docker compose up -d` succeeds and you're happy
# with the setup. It:
#
#   1. Records the exact image digest (SHA256) for each container
#   2. Exports each image as a self-contained .tar file
#   3. Writes a restore.sh you can run on any machine to reload everything
#
# The .tar files are large (~8-15 GB total). Store them somewhere safe
# (external drive, NAS, etc). They never expire and need no internet to load.
#
# Usage:
#   chmod +x archive.sh   (once)
#   ./archive.sh

set -euo pipefail
cd "$(dirname "$0")"

ARCHIVE_DIR="./archive"
mkdir -p "$ARCHIVE_DIR"

SERVICES=("forge" "invokeai")

echo "==> Recording image digests and exporting images..."
echo ""

# Start restore script
RESTORE="$ARCHIVE_DIR/restore.sh"
cat > "$RESTORE" <<'HEADER'
#!/usr/bin/env bash
# Restore saved images to Docker on any machine.
# Run from the directory containing this script.
# Requires Docker to be installed (and NVIDIA Container Toolkit for GPU use).
#
# Usage:
#   chmod +x restore.sh
#   ./restore.sh
set -euo pipefail
cd "$(dirname "$0")"
HEADER

for SERVICE in "${SERVICES[@]}"; do
    # Get the image name from the running container
    IMAGE=$(docker inspect --format='{{.Config.Image}}' "$(docker compose ps -q "$SERVICE")")

    # Get the resolved digest (immutable SHA256 — survives tag changes or deletions)
    DIGEST=$(docker inspect --format='{{index .RepoDigests 0}}' "$IMAGE" 2>/dev/null || \
             docker inspect --format='{{.Id}}' "$IMAGE")

    TARFILE="$ARCHIVE_DIR/${SERVICE}.tar"

    echo "  Service:  $SERVICE"
    echo "  Image:    $IMAGE"
    echo "  Digest:   $DIGEST"
    echo "  Saving → $TARFILE"
    docker save "$IMAGE" -o "$TARFILE"
    echo "  Done."
    echo ""

    # Append load step to restore script
    cat >> "$RESTORE" <<ENTRY

echo "==> Loading $SERVICE image from ${SERVICE}.tar ..."
docker load -i "${SERVICE}.tar"
echo "    Loaded: $IMAGE"
echo "    Digest: $DIGEST"
ENTRY
done

# Finish restore script
cat >> "$RESTORE" <<'FOOTER'

echo ""
echo "==> All images loaded. Copy your docker-compose.yml and data/ directory"
echo "    to this machine, then run: docker compose up -d"
FOOTER

chmod +x "$RESTORE"

# Write a digest record as a plain text reference
DIGEST_FILE="$ARCHIVE_DIR/digests.txt"
echo "# Image digests captured on $(date)" > "$DIGEST_FILE"
for SERVICE in "${SERVICES[@]}"; do
    IMAGE=$(docker inspect --format='{{.Config.Image}}' "$(docker compose ps -q "$SERVICE")")
    DIGEST=$(docker inspect --format='{{index .RepoDigests 0}}' "$IMAGE" 2>/dev/null || \
             docker inspect --format='{{.Id}}' "$IMAGE")
    echo "$SERVICE: $DIGEST" >> "$DIGEST_FILE"
done

echo "==> Archive complete. Files written to $ARCHIVE_DIR/:"
ls -lh "$ARCHIVE_DIR/"
echo ""
echo "    To restore on another machine:"
echo "    1. Copy the archive/ folder and docker-compose.yml to the new machine"
echo "    2. Copy data/ to the new machine (your models and outputs)"
echo "    3. Run: ./archive/restore.sh"
echo "    4. Run: docker compose up -d"