#!/bin/sh
# bty-grow-images-partition - extend the BTY_IMAGES exFAT partition to
# fill the boot stick on first boot.
#
# The cooked usb-iso ISO ships with a pre-allocated 4 GiB exFAT
# partition labelled BTY_IMAGES at the trailing edge of the hybrid
# ISO. After ``dd`` to a stick larger than the artifact, everything
# past the original 4 GiB is unallocated. This script claims that
# tail by deleting + recreating partition #BTY_IMAGES at full disk
# size and re-running mkfs.exfat.
#
# The hybrid ISO is MBR-only (live-build's ``--bootloaders
# syslinux,grub-efi`` produces an isohdpfx-style hybrid, not GPT),
# so this script uses sfdisk throughout. MBR has no secondary
# partition table to relocate.
#
# Driven by ``bty-grow-images-partition.service``, runs once before
# ``bty-tui-on-tty1.service`` and ``bty-flash-on-boot.service`` so
# the grown partition is in place by the time the TUI scans for
# images.
#
# Safety gates (any failure exits 0; this is best-effort):
#
#   1. The script we booted from must be a real block device (skip
#      Ventoy / loopback / netboot - their boot medium is a file
#      inside another fs, growing the host stick would be wrong).
#   2. The BTY_IMAGES partition must live on the same physical disk
#      as the live medium (skip if BTY_IMAGES is on a separate
#      operator-supplied stick - that's their data, not ours to
#      reformat).
#   3. The partition must already be at less than the disk's last
#      usable sector (skip if already grown - idempotent across
#      reboots).
#   4. The partition must be empty (no ``*.img.zst`` / ``*.qcow2`` /
#      ``*.img`` files present - skip if the operator has already
#      dropped images, since exFAT online-grow isn't available on
#      Debian and we'd lose their data).

set -eu

log() {
    printf 'bty-grow-images-partition: %s\n' "$*"
}

# --- 1. Find the live medium's block device ---------------------------------
#
# live-boot mounts the medium at /run/live/medium. ``findmnt`` reads
# the source from /proc/self/mountinfo and gives us /dev/sdXN. If
# /run/live/medium isn't a mountpoint (netboot - kernel/initrd loaded
# over HTTP, no boot medium block device), bail.

medium_src=$(findmnt -no SOURCE /run/live/medium 2>/dev/null || true)
if [ -z "$medium_src" ]; then
    log "no live medium mounted at /run/live/medium; nothing to grow"
    exit 0
fi
if [ ! -b "$medium_src" ]; then
    log "live medium ($medium_src) is not a block device (Ventoy / loopback?); not growing"
    exit 0
fi

medium_disk_name=$(lsblk -no PKNAME "$medium_src" 2>/dev/null || true)
if [ -z "$medium_disk_name" ]; then
    log "could not resolve parent disk of $medium_src; not growing"
    exit 0
fi
medium_disk="/dev/${medium_disk_name}"
log "live medium $medium_src is on disk $medium_disk"

# --- 2. Locate BTY_IMAGES on the same disk ----------------------------------

images_dev=$(blkid -L BTY_IMAGES 2>/dev/null || true)
if [ -z "$images_dev" ]; then
    log "no partition labelled BTY_IMAGES; nothing to grow"
    exit 0
fi
images_disk_name=$(lsblk -no PKNAME "$images_dev" 2>/dev/null || true)
if [ "$images_disk_name" != "$medium_disk_name" ]; then
    log "BTY_IMAGES ($images_dev) is on $images_disk_name, not boot disk $medium_disk_name; not growing"
    exit 0
fi
log "BTY_IMAGES is at $images_dev (same disk as live medium)"

# --- 3. Pull the partition number, current end, and disk end ----------------
#
# blkid resolved the partition to a device path like /dev/sdb3.
# Extract the trailing digits as the partition number; same approach
# handles /dev/nvme0n1p3 (extracts "3") and /dev/sdb3 (extracts "3").
part_num=$(printf '%s\n' "$images_dev" | sed -E 's/^.*[^0-9]([0-9]+)$/\1/')
if ! [ "$part_num" -gt 0 ] 2>/dev/null; then
    log "could not extract partition number from $images_dev; not growing"
    exit 0
fi

# Disk size in 512-byte sectors. The current end-sector of the
# partition comes from sfdisk --json (start + size - 1).
disk_size=$(blockdev --getsz "$medium_disk")
sfdisk_json=$(sfdisk --json "$medium_disk" 2>/dev/null || true)
if [ -z "$sfdisk_json" ]; then
    log "sfdisk --json $medium_disk failed; not growing"
    exit 0
fi

# The grow script can't take a hard jq dependency (not in our chroot
# package list and live-build images stay lean). Use a tiny python3
# one-liner instead - python3 IS in the chroot (we install bty into
# a venv at /opt/bty/venv). Emit ``<start> <size>`` on a single line
# for partition #${part_num}; empty string if not found.
parsed=$(printf '%s' "$sfdisk_json" | python3 -c '
import json, re, sys
table = json.load(sys.stdin)["partitiontable"]
target = sys.argv[1]
for p in table["partitions"]:
    node = p.get("node", "")
    # Trailing digits only; "x86_64" digits in the path must NOT
    # contribute (would yield e.g. 86643 instead of 3).
    m = re.search(r"(\d+)$", node)
    if m and m.group(1) == target:
        print(p["start"], p["size"])
        sys.exit(0)
sys.exit(1)
' "$part_num" 2>/dev/null || true)
if [ -z "$parsed" ]; then
    log "could not parse partition $part_num from sfdisk JSON; not growing"
    exit 0
fi
# parsed is "<start> <size>"; split via POSIX parameter expansion.
part_start=${parsed%% *}
part_size=${parsed##* }
part_end=$((part_start + part_size - 1))

# Reserve a small slop at the very end (1 MiB = 2048 sectors) -
# matches MBR alignment conventions and avoids corner cases where
# the disk's reported size is one sector short of the actual end.
last_usable=$((disk_size - 2048))
slack=$((last_usable - part_end))
log "disk_size=$disk_size last_usable=$last_usable part_start=$part_start part_end=$part_end slack=$slack"

# Allow a small slack tolerance: a freshly-baked artifact may have
# already taken the file to its full extent (the operator dd'd onto
# a stick exactly the artifact's size). Anything under 1 MiB (2048
# sectors) of unallocated tail isn't worth reformatting.
if [ "$slack" -lt 2048 ]; then
    log "BTY_IMAGES already at (or near) the disk's last usable sector; nothing to grow"
    exit 0
fi

# --- 4. Confirm the partition is empty before destroying it -----------------

mp=$(mktemp -d)
# shellcheck disable=SC2064
trap "umount '$mp' 2>/dev/null || true; rmdir '$mp' 2>/dev/null || true" EXIT
if ! mount -o ro "$images_dev" "$mp" 2>/dev/null; then
    log "could not mount $images_dev RO for emptiness check; not growing"
    exit 0
fi
non_empty=""
for ext in img.zst qcow2 img iso; do
    if find "$mp" -maxdepth 2 -type f -name "*.${ext}" -print -quit 2>/dev/null | grep -q .; then
        non_empty="found *.${ext} files"
        break
    fi
done
umount "$mp" 2>/dev/null || true
rmdir "$mp" 2>/dev/null || true
trap - EXIT

if [ -n "$non_empty" ]; then
    log "BTY_IMAGES is NOT empty ($non_empty); preserving operator data, not growing"
    exit 0
fi

# --- 5. Grow: delete the partition, append a fresh one, mkfs.exfat ----------

log "growing BTY_IMAGES (partition $part_num) on $medium_disk to fill the disk"

# sfdisk --delete removes the partition entry; the freed slot is
# reusable. ``--append`` then drops a fresh entry in the next free
# slot at the default start (next free sector after partition 2)
# spanning the rest of the disk. Spec ``,,07`` -> default start,
# default size, type 0x07 (NTFS/exFAT).
sfdisk --delete "$medium_disk" "$part_num"
echo ",,07" | sfdisk --append "$medium_disk"

partprobe "$medium_disk" 2>/dev/null || true
udevadm settle

# Resolve the new partition's device path. /dev/sdb -> /dev/sdb3,
# but /dev/nvme0n1 -> /dev/nvme0n1p3, and same for /dev/mmcblk*.
case "$medium_disk_name" in
    nvme*|mmcblk*|loop*)
        new_dev="${medium_disk}p${part_num}"
        ;;
    *)
        new_dev="${medium_disk}${part_num}"
        ;;
esac

if [ ! -b "$new_dev" ]; then
    log "expected new partition at $new_dev not present after partprobe; aborting"
    exit 0
fi

mkfs.exfat -L BTY_IMAGES "$new_dev"
log "BTY_IMAGES is now $new_dev, full disk size"
