#!/bin/sh
# bty-image-store-init - prepare a block device as the bty image store.
#
# Usage: bty-image-store-init [--yes] /dev/sdX
#
# Partitions the given device with GPT, formats partition 1 as ext4
# with label ``BTY_IMAGE_STORE``, and mounts it at /var/lib/bty/images
# so the image cache survives appliance reflashes. The bty-server
# appliance image ships an /etc/fstab line matching the same label,
# so a freshly-flashed appliance auto-mounts an already-prepared disk
# on boot.
#
# If /var/lib/bty/images currently has files on the rootfs (the
# typical "appliance just installed, no separate store yet" case),
# they're staged into a temporary directory, copied onto the new
# disk after format, and the staging dir removed -- so the operator
# does not lose images by running this script.
#
# Safety rails:
#   - Refuses to format the rootfs disk.
#   - Refuses to format a device with a currently-mounted partition.
#   - Confirmation prompt unless ``--yes`` / ``-y`` is given.
#
# After success:
#   - The disk auto-mounts on every subsequent boot via /etc/fstab.
#   - bty-web is restarted so it picks up the (now properly mounted)
#     image directory.
#   - Reflashing the bty-server appliance preserves the disk's
#     contents; the new appliance's fstab line matches the same
#     label and re-mounts automatically.

set -eu

DEVICE=
ASSUME_YES=

usage() {
    cat <<'EOF'
Usage: bty-image-store-init [--yes|-y] DEVICE

Prepare DEVICE (e.g. /dev/sdb, /dev/nvme1n1) as the bty image store.

The device is wiped, partitioned (GPT, single partition), formatted as
ext4 with label BTY_IMAGE_STORE, mounted at /var/lib/bty/images, and
will auto-mount on subsequent boots (including after reflashing the
bty-server appliance).

  --yes, -y    skip the confirmation prompt
  -h, --help   show this help and exit
EOF
}

while [ $# -gt 0 ]; do
    case "$1" in
        --yes|-y) ASSUME_YES=1; shift ;;
        --help|-h) usage; exit 0 ;;
        --) shift; break ;;
        -*) echo "bty-image-store-init: unknown flag: $1" >&2; usage >&2; exit 2 ;;
        *) DEVICE="$1"; shift; break ;;
    esac
done

if [ -z "${DEVICE:-}" ]; then
    echo "bty-image-store-init: DEVICE argument required" >&2
    usage >&2
    exit 2
fi

if [ "$(id -u)" -ne 0 ]; then
    echo "bty-image-store-init: must be run as root (try sudo)" >&2
    exit 2
fi

if [ ! -b "$DEVICE" ]; then
    echo "bty-image-store-init: $DEVICE is not a block device" >&2
    exit 2
fi

# Refuse to format the rootfs disk. ``findmnt`` returns the rootfs
# partition (e.g. /dev/sda2); strip the trailing partition number /
# ``pN`` suffix to get the parent disk (/dev/sda or /dev/nvme0n1).
ROOTFS_PART=$(findmnt -no SOURCE / 2>/dev/null || true)
case "$ROOTFS_PART" in
    /dev/nvme*|/dev/mmcblk*)
        ROOTFS_DISK=$(printf '%s' "$ROOTFS_PART" | sed 's/p[0-9]\+$//')
        ;;
    /dev/*)
        ROOTFS_DISK=$(printf '%s' "$ROOTFS_PART" | sed 's/[0-9]\+$//')
        ;;
    *)
        ROOTFS_DISK=
        ;;
esac
if [ -n "$ROOTFS_DISK" ] && [ "$DEVICE" = "$ROOTFS_DISK" ]; then
    echo "bty-image-store-init: $DEVICE is the rootfs disk; refusing to format" >&2
    exit 2
fi

# Refuse if any partition of DEVICE is currently mounted somewhere.
# Reading the column with -no MOUNTPOINT keeps the parser simple;
# whitespace-only lines (unmounted partitions) get filtered out.
MOUNTED=$(lsblk -no MOUNTPOINT "$DEVICE" 2>/dev/null | sed '/^[[:space:]]*$/d' || true)
if [ -n "$MOUNTED" ]; then
    echo "bty-image-store-init: $DEVICE has currently-mounted partitions:" >&2
    printf '  %s\n' $MOUNTED >&2
    echo "Unmount them first, then re-run." >&2
    exit 2
fi

IMAGES_DIR=/var/lib/bty/images

# Detect existing content on the rootfs's IMAGES_DIR that should be
# preserved across the format. This only matters when IMAGES_DIR is
# NOT already a separate mountpoint (i.e. it's still on the rootfs).
STAGE_FILES=0
if [ -d "$IMAGES_DIR" ]; then
    CURRENT_MOUNT=$(findmnt -no SOURCE "$IMAGES_DIR" 2>/dev/null || true)
    if [ -z "$CURRENT_MOUNT" ]; then
        COUNT=$(find "$IMAGES_DIR" -mindepth 1 -maxdepth 1 2>/dev/null | wc -l)
        if [ "$COUNT" -gt 0 ]; then
            STAGE_FILES=$COUNT
        fi
    fi
fi

echo
echo "About to:"
echo "  - WIPE all existing data on $DEVICE"
echo "  - Create a GPT label + single partition spanning the whole disk"
echo "  - Format with ext4, label BTY_IMAGE_STORE"
echo "  - Mount at $IMAGES_DIR"
echo "  - Append a matching line to /etc/fstab so it auto-mounts on boot"
if [ "$STAGE_FILES" -gt 0 ]; then
    echo "  - Copy $STAGE_FILES existing file(s) from $IMAGES_DIR onto the new disk"
fi
echo "  - Restart bty-web.service"
echo

if [ -z "$ASSUME_YES" ]; then
    printf 'Type "yes" to continue: '
    read -r CONFIRM
    if [ "$CONFIRM" != "yes" ]; then
        echo "aborted." >&2
        exit 1
    fi
fi

STAGE_DIR=
if [ "$STAGE_FILES" -gt 0 ]; then
    STAGE_DIR=$(mktemp -d /var/lib/bty/.image-store-stage.XXXXXX)
    echo "Staging $IMAGES_DIR -> $STAGE_DIR ..."
    cp -a "$IMAGES_DIR"/. "$STAGE_DIR"/
fi

echo "Wiping $DEVICE ..."
wipefs -a "$DEVICE" >/dev/null

echo "Partitioning $DEVICE (GPT, single partition) ..."
parted --script "$DEVICE" mklabel gpt
parted --script "$DEVICE" mkpart bty-image-store ext4 0% 100%
partprobe "$DEVICE" 2>/dev/null || true
udevadm settle

# Resolve partition device path. lsblk's first ``part`` row under
# DEVICE is partition 1 regardless of naming convention
# (``/dev/sda1`` vs ``/dev/nvme0n1p1``).
PART=$(lsblk -lno NAME,TYPE "$DEVICE" | awk '$2 == "part" {print "/dev/" $1; exit}')
if [ -z "${PART:-}" ] || [ ! -b "$PART" ]; then
    echo "bty-image-store-init: cannot find partition 1 on $DEVICE after partprobe" >&2
    exit 2
fi

echo "Formatting $PART as ext4 (label BTY_IMAGE_STORE) ..."
mkfs.ext4 -F -L BTY_IMAGE_STORE "$PART" >/dev/null

install -d -o bty -g bty -m 0750 "$IMAGES_DIR"

# Idempotent mount: if a previous attempt left something mounted at
# IMAGES_DIR, unmount before remounting the new disk.
if mountpoint -q "$IMAGES_DIR"; then
    umount "$IMAGES_DIR"
fi
mount LABEL=BTY_IMAGE_STORE "$IMAGES_DIR"
chown bty:bty "$IMAGES_DIR"
chmod 0750 "$IMAGES_DIR"

if [ -n "$STAGE_DIR" ]; then
    echo "Restoring staged files to $IMAGES_DIR ..."
    cp -a "$STAGE_DIR"/. "$IMAGES_DIR"/
    chown -R bty:bty "$IMAGES_DIR"
    rm -rf "$STAGE_DIR"
fi

# Add fstab line if it isn't already there. The bty-server appliance
# image ships this same line via the cloud-init bake, so on a fresh
# install it's already present and this is a no-op. The grep guard
# keeps repeated invocations idempotent.
if ! grep -qE '^LABEL=BTY_IMAGE_STORE[[:space:]]' /etc/fstab; then
    echo "Adding fstab line ..."
    printf 'LABEL=BTY_IMAGE_STORE %s ext4 nofail,x-systemd.device-timeout=10s 0 2\n' \
        "$IMAGES_DIR" >> /etc/fstab
    systemctl daemon-reload
fi

echo "Restarting bty-web.service ..."
systemctl restart bty-web.service

echo
echo "Done. $DEVICE is now the bty image store, mounted at $IMAGES_DIR."
echo "It will auto-mount on every boot via LABEL=BTY_IMAGE_STORE,"
echo "including after the bty-server appliance is reflashed."
