#!/usr/bin/env bash
# ============================================================================
# blkstage.sh  —  DriveStage Engine (BlkStage)
# ----------------------------------------------------------------------------
# A Linux-native multi-payload USB toolkit. Inspired by Easy2Boot and RMPrepUSB
# (Steve Si, https://easy2boot.xyz); implemented from scratch for the ext4
# and UEFI era. No code derived from those projects.
#
# Layout created on the target drive:
#   [ Drive: /dev/sdX or /dev/nvmeXn1 ]
#   ├─ Partition 1: ESP            (FAT32, ~512MB)  -> /boot/efi   (UEFI GRUB)
#   ├─ Partition 2: BIOS Boot      (unformatted, 1MB)              (GRUB GPT)
#   └─ Partition 3: Payloads       (ext4, remainder) -> /payloads  (ISOs/tars)
#
# Subcommands:
#   setup    <device>          Partition + format + install GRUB + create dirs
#   deploy   <device> <file>   Copy a .iso / .tar / .img payload (extract tarballs)
#   scan     <device>          Extract tarballs + refresh auto-scan grub.cfg UUID
#   list     <device>          Show registered payloads and sizes
#   shell    <device>          Mount partitions and open a helper shell
#   clean    <device>          Wipe partition table (DESTRUCTIVE)
#
# Usage:
#   sudo ./blkstage.sh setup /dev/sdX
#   sudo ./blkstage.sh deploy /dev/sdX ./ubuntu-24.04.iso
#   sudo ./blkstage.sh scan   /dev/sdX
#
# Requires: bash 4+, parted, sfdisk, mkfs.fat, mkfs.ext4, grub2/grub-efi,
#           rsync, tar, blkid, findmnt
#
# Author: Jeremy Anderson <info@dcos.net>
# License: MIT
# ============================================================================

set -euo pipefail
shopt -s nullglob
shopt -s extglob

# ----------------------------------------------------------------------------
# Configuration (override via environment or /etc/drivestage.conf)
# ----------------------------------------------------------------------------
ESP_SIZE_MB="${ESP_SIZE_MB:-512}"
BIOS_BOOT_SIZE_MB="${BIOS_BOOT_SIZE_MB:-1}"
PAYLOAD_LABEL="${PAYLOAD_LABEL:-DRIVESTAGE}"
ESP_LABEL="${ESP_LABEL:-DRIVESTAGE-EFI}"
MOUNT_BASE="${MOUNT_BASE:-/mnt/drivestage}"
CONF_FILE="${CONF_FILE:-/etc/drivestage.conf}"

# Load optional config (with ownership/mode safety check)
if [[ -r "$CONF_FILE" ]]; then
    _owner="$(stat -c %u "$CONF_FILE" 2>/dev/null || echo 1)"
    _mode="$(stat -c %a "$CONF_FILE" 2>/dev/null || echo 777)"
    if [[ "$_owner" -eq 0 && "$_mode" =~ ^6[04][0-9]$ ]]; then
        # shellcheck disable=SC1090  # CONF_FILE is intentionally dynamic; validated above.
        source "$CONF_FILE"
    else
        echo "[!] Skipping unsafe config file (owner=$_owner mode=$_mode): $CONF_FILE" >&2
    fi
    unset _owner _mode
fi

# ----------------------------------------------------------------------------
# Logging / pretty output
# ----------------------------------------------------------------------------
if [[ -t 1 ]]; then
    C_RED=$'\033[0;31m'; C_GREEN=$'\033[0;32m'; C_YELLOW=$'\033[1;33m'
    C_BLUE=$'\033[0;34m'; C_BOLD=$'\033[1m';    C_DIM=$'\033[2m';  C_NC=$'\033[0m'
else
    C_RED=""; C_GREEN=""; C_YELLOW=""; C_BLUE=""; C_BOLD=""; C_DIM=""; C_NC=""
fi

log_info()  { printf '%s[+]%s %s\n' "$C_GREEN" "$C_NC" "$*"; }
log_warn()  { printf '%s[!]%s %s\n' "$C_YELLOW" "$C_NC" "$*"; }
log_error() { printf '%s[-]%s %s\n' "$C_RED" "$C_NC" "$*" >&2; }
log_step()  { printf '%s[*]%s %s\n' "$C_BLUE" "$C_NC" "$*"; }
log_dim()   { printf '%s%s%s\n' "$C_DIM" "$*" "$C_NC"; }

die()  { log_error "$*"; exit 1; }
abort(){ log_error "$*"; exit 130; }

# ----------------------------------------------------------------------------
# Pre-flight checks
# ----------------------------------------------------------------------------
check_root() {
    [[ $EUID -eq 0 ]] || die "Must run as root (try: sudo $0 $*)."
}

check_commands() {
    # Required utilities — all standard Linux tooling. Anything missing
    # produces a single consolidated error with per-distro install hints.
    local -a required=(
        sfdisk parted mkfs.fat mkfs.ext4 blkid findmnt tar rsync
        wipefs blockdev partprobe udevadm
    )
    local missing=()
    local cmd
    for cmd in "${required[@]}"; do
        command -v "$cmd" >/dev/null 2>&1 || missing+=("$cmd")
    done

    # GRUB install binary has two names across distros — accept either.
    # Step-down resolution in grub_install_cmd() picks the right one at use time.
    command -v grub-install  >/dev/null 2>&1 \
        || command -v grub2-install >/dev/null 2>&1 \
        || missing+=("grub-install (or grub2-install on Fedora)")

    (( ${#missing[@]} )) || return 0
    die "Missing required commands: ${missing[*]}\n" \
        "  Debian/Ubuntu: apt install parted dosfstools e2fsprogs grub-pc-bin grub-efi-amd64-bin util-linux rsync tar\n" \
        "  Fedora:        dnf install parted dosfstools e2fsprogs grub2-pc grub2-efi-x64 grub2-tools util-linux rsync tar\n" \
        "  Arch:          pacman -S parted dosfstools e2fsprogs grub efibootmgr util-linux rsync tar"
}

# Resolve the correct GRUB install binary name for this system.
# Step-down order: Debian/Ubuntu name first, then Fedora/RHEL, then a
# deterministic default that produces a clear error if neither is present.
grub_install_cmd() {
    command -v grub-install  >/dev/null 2>&1 && { printf 'grub-install';  return; }
    command -v grub2-install >/dev/null 2>&1 && { printf 'grub2-install'; return; }
    printf 'grub-install'   # let it fail with a clear error
}

# ----------------------------------------------------------------------------
# Disk helpers
# ----------------------------------------------------------------------------

# Resolve a device argument to a canonical /dev/XXX block device path.
# Resolution order (step-down, Unix philosophy):
#   1. Accept /dev/* paths as-is.
#   2. Accept /disk/by-{id,uuid,path,partlabel} symlinks via readlink.
#   3. Accept bare names (sdb, nvme0n1) by prepending /dev/.
#   4. Verify the result is a block device; abort otherwise.
resolve_device() {
    local dev="$1"
    case "$dev" in
        /dev/*)  : ;;                          # already canonical
        */dev/*) : ;;                          # disk/by-id/... symlink
        *)       dev="/dev/$dev" ;;             # bare name
    esac
    dev="$(readlink -f "$dev")"
    [[ -b "$dev" ]] || die "Not a block device: $dev"
    printf '%s' "$dev"
}

# Is this a whole disk (not a partition)?
is_whole_disk() {
    local dev="$1"
    # /dev/sda, /dev/nvme0n1, /dev/mmcblk0 are whole disks
    # /dev/sda1, /dev/nvme0n1p2 are partitions
    [[ "$dev" =~ ^/dev/(sd[a-z]+|nvme[0-9]+n[0-9]+|mmcblk[0-9]+|vd[a-z]+|loop[0-9]+)$ ]]
}

# Get the partition node naming scheme for a disk.
# /dev/sda      -> /dev/sda1, /dev/sda2  (suffix is just the number)
# /dev/nvme0n1  -> /dev/nvme0n1p1        (suffix is p<number>)
# /dev/mmcblk0  -> /dev/mmcblk0p1        (suffix is p<number>)
part_prefix() {
    local disk="$1"
    case "$disk" in
        /dev/nvme*|/dev/mmcblk*|/dev/loop*) printf '%sp' "$disk" ;;
        *)                                   printf '%s'  "$disk" ;;
    esac
}

# Total size of a block device in megabytes (integer).
disk_size_mb() {
    local disk="$1"
    local bytes
    bytes="$(blockdev --getsize64 "$disk")"
    echo $(( bytes / 1024 / 1024 ))
}

# List candidate removable disks for interactive selection.
list_removable_disks() {
    local name path size_gb
    log_step "Detecting removable/USB candidate disks:"
    printf '  %-12s %-10s %-12s %s\n' "DEVICE" "SIZE" "TRANSPORT" "MODEL"
    printf '  %-12s %-10s %-12s %s\n' "------" "----" "---------" "-----"
    for path in /sys/block/*; do
        name="$(basename "$path")"
        case "$name" in
            loop*|ram*|sr*) continue ;;
        esac
        # Read transport / model / size directly. `cat`/`ls` over /sys is
        # safe here — entries are kernel-controlled and have known shapes.
        local rm tran model size_bytes
        rm="$(cat "$path/removable" 2>/dev/null || echo 0)"
        tran="$(readlink -f "$path/device" 2>/dev/null | sed 's|.*/||' || echo '?')"
        model="$(tr -d ' ' < "$path/device/model" 2>/dev/null || echo '?')"
        size_bytes="$(blockdev --getsize64 "/dev/$name" 2>/dev/null || echo 0)"
        size_gb="$(awk "BEGIN{printf \"%.1f\", $size_bytes/1024/1024/1024}")"
        # Show if removable OR explicitly USB transport
        if [[ "$rm" == 1 || "$tran" == *usb* || "$tran" == *mmc* ]]; then
            printf '  %-12s %-10s %-12s %s\n' "/dev/$name" "${size_gb}GB" "${tran:-?}" "${model:-?}"
        fi
    done
}

# Interactively pick a disk if none provided.
pick_disk_interactive() {
    list_removable_disks
    echo
    read -rp "Enter device to use (e.g. /dev/sdb): " dev
    [[ -n "$dev" ]] || die "No device selected."
    resolve_device "$dev"
}

# ----------------------------------------------------------------------------
# Safety guards
# ----------------------------------------------------------------------------

# Refuse to operate on the running OS disk.
# Detects plain partitions (/dev/sda3), LVM (/dev/mapper/vg-root),
# LUKS (/dev/mapper/cryptroot), mdadm (/dev/md0), and any device-mapper
# descendant by walking lsblk's parent (PKNAME) chain from each mounted
# SOURCE back to a whole-disk ancestor.
#
# Critical-mount check uses a single case statement (step-down) rather
# than a for-loop+if: pattern lists in case are the POSIX-native way to
# express "any of these prefixes match."
is_critical_mount() {
    case "$1" in
        /|/boot|/boot/efi|/usr|/var|/home|/etc|/efi|/opt|/srv) return 0 ;;
        /boot/*|/usr/*|/var/*|/home/*|/etc/*|/efi/*|/opt/*|/srv/*) return 0 ;;
        *)                                                          return 1 ;;
    esac
}

guard_root_disk() {
    local disk="$1"
    local mnt dev src holder

    while read -r mnt dev _; do
        is_critical_mount "$mnt" || continue

        # Walk the holder chain: dev -> PKNAME -> PKNAME -> ... -> whole disk.
        src="$dev"
        while [[ -n "$src" ]]; do
            # Direct match: dev is $disk itself, or a partition of $disk
            # (e.g. /dev/sda3 is a partition of /dev/sda).
            case "$src" in
                "$disk"|"$disk"?*) die "Refusing: '$disk' backs mounted filesystem '$mnt' (likely your OS disk)." ;;
            esac
            # Resolve device-mapper / LVM / LUKS / mdadm ancestry via lsblk PKNAME.
            holder="$(lsblk -ndo PKNAME "$src" 2>/dev/null || true)"
            [[ "$holder" == "$src" ]] && break   # prevent infinite loop
            src="$holder"
        done
    done < <(findmnt -nno TARGET,SOURCE)
}

confirm_destructive() {
    local disk="$1"
    local action="${2:-format}"
    echo
    log_warn "About to $action device: $disk"
    log_warn "ALL DATA ON $disk WILL BE LOST."
    local size_gb
    size_gb="$(awk "BEGIN{printf \"%.1f\", $(blockdev --getsize64 "$disk")/1024/1024/1024}")"
    printf '  Device : %s\n' "$disk"
    printf '  Size   : %s GB\n' "$size_gb"
    printf '  Model  : %s\n' "$(lsblk -ndo MODEL "$disk" 2>/dev/null || echo '?')"
    echo
    read -rp "Type the device name (e.g. $(basename "$disk")) to confirm: " confirm
    [[ "$confirm" == "$(basename "$disk")" ]] || abort "Aborted (no match)."
}

# ----------------------------------------------------------------------------
# Partitioning
# ----------------------------------------------------------------------------

wipe_signatures() {
    local disk="$1"
    log_step "Wiping existing signatures on $disk..."
    wipefs -af "$disk" 2>/dev/null || die "wipefs failed on $disk (disk may be read-only or in use)"
    dd if=/dev/zero of="$disk" bs=1M count=1 status=none conv=notrunc 2>/dev/null || true
    partprobe "$disk" 2>/dev/null || true
    udevadm settle 2>/dev/null || true
}

# Create the 3-partition GPT layout using sfdisk.
#   1: ESP            FAT32,  EFI System Partition type (EF00)
#   2: BIOS Boot      1 MiB,  BIOS boot type        (EF02)
#   3: Payloads       ext4,   Linux filesystem type (8300)
create_partitions() {
    local disk="$1"
    local total_mb
    total_mb="$(disk_size_mb "$disk")"

    log_step "Creating GPT partition table on $disk (total ${total_mb} MiB)..."
    sfdisk "$disk" <<EOF
label: gpt
unit: sectors
first-lba: 2048

# ESP — UEFI System Partition (FAT32)
start=2048, size=$(( ESP_SIZE_MB * 1024 * 1024 / 512 )), type=EF00, name=ESP

# BIOS Boot Partition (for GRUB i386-pc embedding, unformatted)
start=$(( 2048 + ESP_SIZE_MB * 1024 * 1024 / 512 )), size=$(( BIOS_BOOT_SIZE_MB * 1024 * 1024 / 512 )), type=EF02, name=BIOSBOOT

# Payload Partition (ext4, holds ISOs / rootfs tarballs / images)
start=$(( 2048 + (ESP_SIZE_MB + BIOS_BOOT_SIZE_MB) * 1024 * 1024 / 512 )), type=8300, name=PAYLOADS
EOF

    # Re-read partition table
    partprobe "$disk" 2>/dev/null || true
    udevadm settle 2>/dev/null || true
    sleep 1
    log_info "Partition table created."
    sfdisk -l "$disk" 2>/dev/null | sed 's/^/    /'
}

# Get the partition device node for slot N on a disk (waits for udev).
get_partition() {
    local disk="$1" n="$2"
    local prefix
    prefix="$(part_prefix "$disk")"
    local p="${prefix}${n}"
    # Wait up to 30s for udev to materialize the node
    local tries=0
    while [[ ! -b "$p" && $tries -lt 300 ]]; do
        sleep 0.1
        tries=$((tries + 1))
    done
    [[ -b "$p" ]] || die "Partition $p not found after 30s. Check dmesg or run 'partprobe $disk' manually."
    printf '%s' "$p"
}

format_partitions() {
    local disk="$1"
    local p1 p3
    p1="$(get_partition "$disk" 1)"
    p3="$(get_partition "$disk" 3)"

    log_step "Formatting ESP ($p1) as FAT32..."
    mkfs.fat -F 32 -n "$ESP_LABEL" "$p1"

    log_step "Formatting payload partition ($p3) as ext4..."
    # -F = force (in case of stale signature), -L = label, -m 1 = 1% reserved blocks
    mkfs.ext4 -F -L "$PAYLOAD_LABEL" -m 1 -O ^64bit "$p3" 2>&1 | sed 's/^/    /'

    log_info "Filesystems created."
}

# ----------------------------------------------------------------------------
# Mounting
# ----------------------------------------------------------------------------

mount_payload_partition() {
    local disk="$1"
    local p3 mnt
    p3="$(get_partition "$disk" 3)"
    mnt="${MOUNT_BASE}/payload"
    mkdir -p "$mnt"
    if findmnt -n "$p3" >/dev/null; then
        log_dim "    $p3 already mounted at $(findmnt -n -o TARGET "$p3")"
    else
        log_step "Mounting $p3 at $mnt..."
        mount "$p3" "$mnt"
    fi
    printf '%s' "$mnt"
}

mount_esp_partition() {
    local disk="$1" payload_mnt="$2"
    local p1 mnt
    p1="$(get_partition "$disk" 1)"
    mnt="${payload_mnt}/boot/efi"
    mkdir -p "$mnt"
    if findmnt -n "$p1" >/dev/null; then
        log_dim "    $p1 already mounted at $(findmnt -n -o TARGET "$p1")"
    else
        log_step "Mounting $p1 at $mnt..."
        mount "$p1" "$mnt"
    fi
    printf '%s' "$mnt"
}

mount_all() {
    local disk="$1"
    local payload_mnt esp_mnt
    payload_mnt="$(mount_payload_partition "$disk")"
    esp_mnt="$(mount_esp_partition "$disk" "$payload_mnt")"
    printf '%s\n%s\n' "$payload_mnt" "$esp_mnt"
}

unmount_all() {
    local disk="$1"
    log_step "Unmounting partitions on $disk..."
    # Primary: recursive unmount of anything under MOUNT_BASE (handles nested ESP).
    # Step-down: try umount -R; if it fails or MOUNT_BASE is missing, fall through
    # to per-partition cleanup. `|| true` keeps the line from aborting under
    # `set -e` when umount -R encounters a busy mount.
    if [[ -d "$MOUNT_BASE" ]]; then
        umount -R "$MOUNT_BASE" 2>/dev/null || true
    fi

    # Step-down per-partition cleanup: deepest mount first.
    # Resolves any stragglers that umount -R missed (rare, but seen on
    # busy loop devices during test_loop_device.sh).
    local p
    for p in \
        "$(get_partition "$disk" 1 2>/dev/null || true)" \
        "$(get_partition "$disk" 3 2>/dev/null || true)"; do
        [[ -n "$p" ]] || continue
        findmnt -n -o TARGET --source "$p" 2>/dev/null \
            | tac \
            | while read -r m; do umount "$m" 2>/dev/null || true; done
    done
    sync
}

# ----------------------------------------------------------------------------
# Directory structure
# ----------------------------------------------------------------------------

create_directory_structure() {
    local payload_mnt="$1"
    log_step "Creating payload directory structure..."
    local -a dirs=(
        "$payload_mnt/payloads"
        "$payload_mnt/payloads/isos"
        "$payload_mnt/payloads/rootfs_tarballs"
        "$payload_mnt/payloads/rootfs"
        "$payload_mnt/payloads/images"
        "$payload_mnt/payloads/overlay"
        "$payload_mnt/boot/grub"
        "$payload_mnt/boot/grub/conf.d"
        "$payload_mnt/boot/grub/templates"
        "$payload_mnt/boot/efi"
    )
    for d in "${dirs[@]}"; do mkdir -p "$d"; done

    # README drop-zone
    cat > "$payload_mnt/payloads/README.md" <<'EOF'
# DriveStage — Payload Drop Zone

**Just drop files in and boot. No scan command needed.**
GRUB auto-discovers payloads at boot time.

## payloads/isos/          ← Drop .iso files here
Live ISOs are auto-detected at boot. Distro-specific kernel arguments are
applied based on filename pattern matching (ubuntu, arch, debian, fedora,
manjaro, opensuse, kali, alpine, and 20+ more). Unrecognized ISOs fall back
to GRUB's loopback.cfg chainload.

Examples:
  ubuntu-24.04-desktop-amd64.iso
  archlinux-2026.08.01-x86_64.iso
  debian-live-12.5.0-amd64-gnome.iso
  fedora-workstation-live-x86_64-40.iso

## payloads/rootfs/         ← Drop pre-extracted rootfs directories here
Any subdirectory containing /boot/vmlinuz* or /boot/bzImage is auto-discovered.
Kernel + initrd are loaded directly from the ext4 filesystem (no loopback needed).

## payloads/rootfs_tarballs/  ← Drop .tar(.gz/.xz/.zst) here
Tarballs CANNOT be booted directly — GRUB has no untar capability at boot time.
Run `./blkstage.sh scan /dev/sdX` once on your host to extract them into
payloads/rootfs/<name>/, after which they're auto-discovered.
Or use `./blkstage.sh deploy /dev/sdX rootfs.tar.gz` which extracts on copy.

## payloads/images/         ← Drop .img/.raw files here (experimental)
GRUB cannot chainload raw disk images. These appear as menu entries with an
informational message. Convert to ISO or extract the kernel+initrd instead.

## payloads/overlay/        ← (reserved for future overlayfs persistence)

## Adding custom menu entries
Edit /boot/grub/custom.cfg on this partition. It's sourced at boot time
after the auto-scan, so your custom entries appear at the bottom of the menu.

## Workflow summary
1. Drop .iso files into payloads/isos/
2. Eject the USB
3. Boot from it
4. GRUB auto-scans and presents a menu
EOF
    log_info "Directory structure ready."
}

# ----------------------------------------------------------------------------
# GRUB installation
# ----------------------------------------------------------------------------

install_grub() {
    local disk="$1" payload_mnt="$2" esp_mnt="$3"
    local grub_failed=0
    local grub_bin
    grub_bin="$(grub_install_cmd)"

    log_step "Installing GRUB for UEFI (x86_64-efi)..."
    "$grub_bin" \
        --target=x86_64-efi \
        --efi-directory="$esp_mnt" \
        --boot-directory="${payload_mnt}/boot" \
        --removable \
        --recheck 2>&1 | sed 's/^/    /' || grub_failed=1
    if [[ "$grub_failed" -ne 0 ]]; then
        die "UEFI GRUB install failed. Aborting setup to prevent an unbootable drive."
    fi

    log_step "Installing GRUB for BIOS (i386-pc)..."
    "$grub_bin" \
        --target=i386-pc \
        --boot-directory="${payload_mnt}/boot" \
        --recheck \
        "$disk" 2>&1 | sed 's/^/    /' || grub_failed=1
    if [[ "$grub_failed" -ne 0 ]]; then
        die "BIOS GRUB install failed. Aborting setup to prevent an unbootable drive."
    fi

    log_info "GRUB installed for both UEFI and BIOS targets."
}

# Generate the top-level grub.cfg with BOOT-TIME AUTO-SCAN logic.
# No host-side menu generation required — GRUB discovers payloads at boot.
generate_grub_main_cfg() {
    local payload_mnt="$1"
    local p3 p3_uuid
    p3="$(get_partition "$CURRENT_DISK" 3)"
    p3_uuid="$(blkid -s UUID -o value "$p3")"

    local grub_cfg="${payload_mnt}/boot/grub/grub.cfg"
    log_step "Generating auto-scan $grub_cfg..."

    cat > "$grub_cfg" <<'GRUBEOF'
# ============================================================================
# grub.cfg  —  Boot-time Auto-Scan Version
# ----------------------------------------------------------------------------
# This config AUTO-DISCovers payloads at boot time. Just drop files into
# payloads/{isos,rootfs,images}/ on the ext4 partition and reboot — no
# host-side 'scan' command required.
#
# Recognized payload locations:
#   /payloads/isos/*.iso        -> distro-aware loopback boot
#   /payloads/rootfs/*/         -> direct kernel + initrd boot
#   /payloads/images/*.img      -> experimental chainload note
#
# Generated by blkstage.sh
# ============================================================================

# --- Payload partition UUID (baked in at setup time) ---
GRUB_EOF_UUID_PLACEHOLDER

# --- Boot timeout / default ---
set timeout=30
set timeout_style=menu
set default=0

# --- Required GRUB modules ---
insmod part_gpt
insmod ext2
insmod fat
insmod loopback
insmod regexp
insmod all_video
insmod gfxterm
insmod normal
insmod test
insmod search
insmod search_fs_uuid
insmod search_label
insmod configfile

# --- Locate the payload partition by UUID ---
search --no-floppy --fs-uuid --set=root $ds_payload_uuid

# ============================================================================
# Helper: emit the correct ISO menuentry based on filename pattern matching.
# Args: $1 = iso_path (full path), $2 = iso_name (basename)
# ============================================================================
function ds_iso_entry {
    set iso_path="$1"
    set iso_name="$2"
    set distro="unknown"

    # --- Distro detection via regexp (POSIX BRE, case-sensitive) ---
    # Most real-world ISO filenames are lowercase, so we match lowercase.
    if regexp --quiet --ignore-case 'ubuntu' "$iso_name"; then set distro="ubuntu"
    elif regexp --quiet --ignore-case 'archlinux' "$iso_name"; then set distro="arch"
    elif regexp --quiet --ignore-case 'arch-' "$iso_name"; then set distro="arch"
    elif regexp --quiet --ignore-case 'debian.*live' "$iso_name"; then set distro="debian-live"
    elif regexp --quiet --ignore-case 'debian-live' "$iso_name"; then set distro="debian-live"
    elif regexp --quiet --ignore-case 'debian' "$iso_name"; then set distro="debian"
    elif regexp --quiet --ignore-case 'manjaro' "$iso_name"; then set distro="manjaro"
    elif regexp --quiet --ignore-case 'fedora' "$iso_name"; then set distro="fedora"
    elif regexp --quiet --ignore-case 'centos-stream' "$iso_name"; then set distro="centos_stream"
    elif regexp --quiet --ignore-case 'centos_stream' "$iso_name"; then set distro="centos_stream"
    elif regexp --quiet --ignore-case 'centos' "$iso_name"; then set distro="centos"
    elif regexp --quiet --ignore-case 'rocky' "$iso_name"; then set distro="rocky"
    elif regexp --quiet --ignore-case 'alma' "$iso_name"; then set distro="alma"
    elif regexp --quiet --ignore-case 'tumbleweed' "$iso_name"; then set distro="opensuse"
    elif regexp --quiet --ignore-case 'leap' "$iso_name"; then set distro="opensuse"
    elif regexp --quiet --ignore-case 'opensuse' "$iso_name"; then set distro="opensuse"
    elif regexp --quiet --ignore-case 'kali' "$iso_name"; then set distro="kali"
    elif regexp --quiet --ignore-case 'parrot' "$iso_name"; then set distro="parrot"
    elif regexp --quiet --ignore-case 'tails' "$iso_name"; then set distro="tails"
    elif regexp --quiet --ignore-case 'knoppix' "$iso_name"; then set distro="knoppix"
    elif regexp --quiet --ignore-case 'systemrescue' "$iso_name"; then set distro="systemrescue"
    elif regexp --quiet --ignore-case 'system-rescue' "$iso_name"; then set distro="systemrescue"
    elif regexp --quiet --ignore-case 'alpine' "$iso_name"; then set distro="alpine"
    elif regexp --quiet --ignore-case 'void' "$iso_name"; then set distro="void"
    elif regexp --quiet --ignore-case 'gentoo' "$iso_name"; then set distro="gentoo"
    elif regexp --quiet --ignore-case 'slax' "$iso_name"; then set distro="slax"
    elif regexp --quiet --ignore-case 'tinycore' "$iso_name"; then set distro="tinycore"
    elif regexp --quiet --ignore-case 'clonezilla' "$iso_name"; then set distro="clonezilla"
    elif regexp --quiet --ignore-case 'gparted' "$iso_name"; then set distro="gparted"
    elif regexp --quiet --ignore-case 'mint' "$iso_name"; then set distro="mint"
    elif regexp --quiet --ignore-case 'elementary' "$iso_name"; then set distro="elementary"
    elif regexp --quiet --ignore-case 'pop-os' "$iso_name"; then set distro="popos"
    elif regexp --quiet --ignore-case 'pop_os' "$iso_name"; then set distro="popos"
    fi

    # --- Casper family (Ubuntu, Mint, elementary, Pop!_OS, Kali, Parrot, Tails) ---
    if [ "$distro" = "ubuntu" ] || [ "$distro" = "mint" ] || [ "$distro" = "elementary" ] || [ "$distro" = "popos" ] || [ "$distro" = "kali" ] || [ "$distro" = "parrot" ] || [ "$distro" = "tails" ]; then
        menuentry "$iso_name  [casper]" "$iso_path" {
            set iso_path="$1"
            search --no-floppy --fs-uuid --set=root $ds_payload_uuid
            loopback loop "$iso_path"
            linux (loop)/casper/vmlinuz boot=casper iso-scan/filename=$iso_path quiet splash ---
            initrd (loop)/casper/initrd
        }
    # --- Live-boot family (Debian Live, Clonezilla, GParted) ---
    elif [ "$distro" = "debian-live" ] || [ "$distro" = "clonezilla" ] || [ "$distro" = "gparted" ]; then
        menuentry "$iso_name  [live-boot]" "$iso_path" {
            set iso_path="$1"
            search --no-floppy --fs-uuid --set=root $ds_payload_uuid
            loopback loop "$iso_path"
            linux (loop)/live/vmlinuz boot=live components findiso=$iso_path quiet splash
            initrd (loop)/live/initrd.img
        }
    # --- Debian netinst / DVD (chainload via loopback.cfg) ---
    elif [ "$distro" = "debian" ]; then
        menuentry "$iso_name  [debian chainload]" "$iso_path" {
            set iso_path="$1"
            search --no-floppy --fs-uuid --set=root $ds_payload_uuid
            loopback loop "$iso_path"
            if [ -f (loop)/boot/grub/loopback.cfg ]; then
                configfile (loop)/boot/grub/loopback.cfg
            else
                echo "No loopback.cfg found in this Debian ISO."
                echo "Manual menuentry may be needed in /boot/grub/custom.cfg"
                sleep 5
            fi
        }
    # --- Arch Linux ---
    elif [ "$distro" = "arch" ]; then
        menuentry "$iso_name  [arch]" "$iso_path" {
            set iso_path="$1"
            search --no-floppy --fs-uuid --set=root $ds_payload_uuid
            loopback loop "$iso_path"
            linux (loop)/arch/boot/x86_64/vmlinuz-linux img_dev=/dev/disk/by-uuid/$ds_payload_uuid img_loop=$iso_path
            initrd (loop)/arch/boot/x86_64/initramfs-linux.img
        }
    # --- Manjaro ---
    elif [ "$distro" = "manjaro" ]; then
        menuentry "$iso_name  [manjaro]" "$iso_path" {
            set iso_path="$1"
            search --no-floppy --fs-uuid --set=root $ds_payload_uuid
            loopback loop "$iso_path"
            linux (loop)/boot/manjaro/x86_64/vmlinuz-x86_64 img_dev=/dev/disk/by-uuid/$ds_payload_uuid img_loop=$iso_path driver=free
            initrd (loop)/boot/manjaro/x86_64/initramfs-x86_64.img
        }
    # --- Fedora ---
    elif [ "$distro" = "fedora" ]; then
        menuentry "$iso_name  [fedora]" "$iso_path" {
            set iso_path="$1"
            search --no-floppy --fs-uuid --set=root $ds_payload_uuid
            loopback loop "$iso_path"
            linux (loop)/isolinux/vmlinuz root=live:CDLABEL=Fedora rd.live.image iso-scan/filename=$iso_path
            initrd (loop)/isolinux/initrd.img
        }
    # --- RHEL family (CentOS Stream, Rocky, Alma) ---
    elif [ "$distro" = "centos_stream" ] || [ "$distro" = "rocky" ] || [ "$distro" = "alma" ]; then
        menuentry "$iso_name  [$distro]" "$iso_path" {
            set iso_path="$1"
            set dname="$2"
            search --no-floppy --fs-uuid --set=root $ds_payload_uuid
            loopback loop "$iso_path"
            linux (loop)/isolinux/vmlinuz root=live:CDLABEL=${dname}-BaseOS-x86_64 rd.live.image iso-scan/filename=$iso_path
            initrd (loop)/isolinux/initrd.img
        }
    # --- CentOS (BIOS-era release) ---
    elif [ "$distro" = "centos" ]; then
        menuentry "$iso_name  [centos]" "$iso_path" {
            set iso_path="$1"
            search --no-floppy --fs-uuid --set=root $ds_payload_uuid
            loopback loop "$iso_path"
            linux (loop)/isolinux/vmlinuz root=live:CDLABEL=CentOS-BaseOS-x86_64 rd.live.image iso-scan/filename=$iso_path
            initrd (loop)/isolinux/initrd.img
        }
    # --- openSUSE ---
    elif [ "$distro" = "opensuse" ]; then
        menuentry "$iso_name  [opensuse]" "$iso_path" {
            set iso_path="$1"
            search --no-floppy --fs-uuid --set=root $ds_payload_uuid
            loopback loop "$iso_path"
            linux (loop)/boot/x86_64/loader/linux isoboot=$iso_path
            initrd (loop)/boot/x86_64/loader/initrd
        }
    # --- SystemRescue ---
    elif [ "$distro" = "systemrescue" ]; then
        menuentry "$iso_name  [systemrescue]" "$iso_path" {
            set iso_path="$1"
            search --no-floppy --fs-uuid --set=root $ds_payload_uuid
            loopback loop "$iso_path"
            linux (loop)/sysresccd/boot/x86_64/vmlinuz img_label=$ds_payload_uuid img_loop=$iso_path archisobasedir=sysresccd
            initrd (loop)/sysresccd/boot/x86_64/sysresccd.img
        }
    # --- Alpine ---
    elif [ "$distro" = "alpine" ]; then
        menuentry "$iso_name  [alpine]" "$iso_path" {
            set iso_path="$1"
            search --no-floppy --fs-uuid --set=root $ds_payload_uuid
            loopback loop "$iso_path"
            linux (loop)/boot/vmlinuz-lts modules=loop,squashfs,sd-mod,usb-storage modloop=$iso_path/boot/modloop-lts alpine_dev=$iso_path iso-scan/filename=$iso_path
            initrd (loop)/boot/initramfs-lts
        }
    # --- Void Linux ---
    elif [ "$distro" = "void" ]; then
        menuentry "$iso_name  [void]" "$iso_path" {
            set iso_path="$1"
            search --no-floppy --fs-uuid --set=root $ds_payload_uuid
            loopback loop "$iso_path"
            linux (loop)/boot/vmlinuz iso-scan/filename=$iso_path root=live:CDLABEL=VOID_LIVE rd.live.image
            initrd (loop)/boot/initrd
        }
    # --- Unknown: fall back to loopback.cfg chainload ---
    else
        menuentry "$iso_name  [auto-chainload]" "$iso_path" {
            set iso_path="$1"
            search --no-floppy --fs-uuid --set=root $ds_payload_uuid
            loopback loop "$iso_path"
            if [ -f (loop)/boot/grub/loopback.cfg ]; then
                configfile (loop)/boot/grub/loopback.cfg
            elif [ -f (loop)/EFI/BOOT/grub.cfg ]; then
                configfile (loop)/EFI/BOOT/grub.cfg
            else
                echo "No loopback.cfg found in $iso_path."
                echo "This ISO needs a manual menuentry."
                echo "Add one to /boot/grub/custom.cfg on the payload partition."
                sleep 5
            fi
        }
    fi
}

# ============================================================================
# Helper: emit a menuentry for a pre-extracted rootfs directory.
# Args: $1 = rootfs_dir (full path), $2 = name (basename)
# ============================================================================
function ds_rootfs_entry {
    set rdir="$1"
    set rname="$2"

    # Detect kernel type and emit the right menuentry
    if [ -f "$rdir/boot/vmlinuz-linux" ]; then
        # Arch-style kernel
        menuentry "$rname  [arch-rootfs]" "$rdir" {
            set rootfs_dir="$1"
            search --no-floppy --fs-uuid --set=root $ds_payload_uuid
            linux ${rootfs_dir}/boot/vmlinuz-linux root=UUID=$ds_payload_uuid rw systemd.post_install=${rootfs_dir}
            initrd ${rootfs_dir}/boot/initramfs-linux.img
        }
    elif [ -f "$rdir/boot/vmlinuz" ]; then
        # Generic vmlinuz
        menuentry "$rname  [rootfs]" "$rdir" {
            set rootfs_dir="$1"
            search --no-floppy --fs-uuid --set=root $ds_payload_uuid
            linux ${rootfs_dir}/boot/vmlinuz root=UUID=$ds_payload_uuid rw init=${rootfs_dir}/sbin/init
            initrd ${rootfs_dir}/boot/initrd.img
        }
    elif [ -f "$rdir/boot/bzImage" ]; then
        # Bare-metal Linux bzImage
        menuentry "$rname  [linux-rootfs]" "$rdir" {
            set rootfs_dir="$1"
            search --no-floppy --fs-uuid --set=root $ds_payload_uuid
            linux ${rootfs_dir}/boot/bzImage root=UUID=$ds_payload_uuid rw init=${rootfs_dir}/sbin/init
        }
    fi
}

# ============================================================================
# Boot-time auto-scan
# ============================================================================

# Clear screen and show banner
clear
echo ""
echo "  ===================================================="
echo "    DriveStage  —  Auto-Scan Payload Menu"
echo "    Partition UUID: $ds_payload_uuid"
echo "  ===================================================="
echo ""

# --- Scan ISOs (case-insensitive extension match) ---
for f in /payloads/isos/*.[iI][sS][oO]; do
    if [ -f "$f" ]; then
        # Try to extract basename via regexp; fall back to full path
        set basename=""
        regexp --set=basename=1 '\([^/]*\)$' "$f"
        if [ -z "$basename" ]; then
            set basename="$f"
        fi
        ds_iso_entry "$f" "$basename"
    fi
done

# --- Scan pre-extracted rootfs directories ---
for d in /payloads/rootfs/*; do
    if [ -d "$d" ]; then
        set rname=""
        regexp --set=rname=1 '\([^/]*\)$' "$d"
        if [ -z "$rname" ]; then
            set rname="$d"
        fi
        ds_rootfs_entry "$d" "$rname"
    fi
done

# --- Scan raw disk images (experimental) ---
for f in /payloads/images/*.[iI][mM][gG]; do
    if [ -f "$f" ]; then
        set basename=""
        regexp --set=basename=1 '\([^/]*\)$' "$f"
        if [ -z "$basename" ]; then
            set basename="$f"
        fi
        menuentry "$basename  [raw image - experimental]" "$f" {
            set img_path="$1"
            search --no-floppy --fs-uuid --set=root $ds_payload_uuid
            echo "Raw disk image booting is experimental."
            echo "Path: $img_path"
            echo ""
            echo "GRUB cannot chainload raw .img files directly."
            echo "Convert to ISO, or extract its kernel+initrd into payloads/rootfs/."
            sleep 10
        }
    fi
done

for f in /payloads/images/*.[rR][aA][wW]; do
    if [ -f "$f" ]; then
        set basename=""
        regexp --set=basename=1 '\([^/]*\)$' "$f"
        if [ -z "$basename" ]; then
            set basename="$f"
        fi
        menuentry "$basename  [raw image - experimental]" "$f" {
            set img_path="$1"
            search --no-floppy --fs-uuid --set=root $ds_payload_uuid
            echo "Raw disk image booting is experimental."
            echo "Path: $img_path"
            sleep 10
        }
    fi
done

# --- Source custom user entries if present ---
if [ -f /boot/grub/custom.cfg ]; then
    source /boot/grub/custom.cfg
fi

# --- Utility entries ---
menuentry "Drop to GRUB command line" {
    configfile (cmdpath)
}

menuentry "Reboot system" {
    reboot
}

menuentry "Halt system" {
    halt
}
GRUBEOF

    # Substitute the UUID placeholder (done outside the heredoc to keep
    # the rest of the GRUB script untouched by bash expansion).
    sed -i "s|^GRUB_EOF_UUID_PLACEHOLDER$|set ds_payload_uuid=\"${p3_uuid}\"|" "$grub_cfg"

    log_info "Auto-scan grub.cfg written. Drop ISOs into payloads/isos/ and reboot."
}

# ----------------------------------------------------------------------------
# Distro detection — bash-side pure function for testing
# ----------------------------------------------------------------------------
# This function mirrors the GRUB-side regexp cascade in ds_iso_entry (inside
# generate_grub_main_cfg's heredoc). It exists as a testable bash interface —
# the GRUB auto-scan uses its own inline matching since GRUB cannot call bash
# functions. Keep both in sync when adding distros.
# Returns "unknown" if no pattern matches.
detect_distro_profile() {
    local filename="$1" lower
    lower="$(printf '%s' "$filename" | tr '[:upper:]' '[:lower:]')"

    case "$lower" in
        *ubuntu*)                                printf 'ubuntu' ;;
        *debian*live*|*debian*netinst*) printf 'debian' ;;
        *debian*)                                printf 'debian' ;;
        *archlinux*|*arch-*)                     printf 'arch' ;;
        *manjaro*)                               printf 'manjaro' ;;
        *fedora*)                                printf 'fedora' ;;
        *centos*stream*)                       printf 'centos_stream' ;;
        *centos*)                                printf 'centos' ;;
        *rocky*)                                 printf 'rocky' ;;
        *alma*)                                  printf 'alma' ;;
        *opensuse*|*tumbleweed*|*leap*)          printf 'opensuse' ;;
        *kali*)                                  printf 'kali' ;;
        *parrot*)                                printf 'parrot' ;;
        *tails*)                                 printf 'tails' ;;
        *knoppix*)                               printf 'knoppix' ;;
        *systemrescue*|*system-rescue*)          printf 'systemrescue' ;;
        *alpine*)                                printf 'alpine' ;;
        *void*)                                  printf 'void' ;;
        *gentoo*)                                printf 'gentoo' ;;
        *slax*)                                  printf 'slax' ;;
        *tinycore*)                              printf 'tinycore' ;;
        *clonezilla*)                            printf 'clonezilla' ;;
        *gparted*)                               printf 'gparted' ;;
        *mint*.iso)                              printf 'mint' ;;
        *elementary*.iso)                        printf 'elementary' ;;
        *pop_os*|*pop-os*)                       printf 'popos' ;;
        *)                                       printf 'unknown' ;;
    esac
}

# ----------------------------------------------------------------------------
# Payload scanner / config refresher
# ----------------------------------------------------------------------------
# NOTE: With the boot-time auto-scan grub.cfg, this function no longer
# generates per-file menuentries. GRUB discovers payloads at boot time.
# This function now:
#   1. Extracts any unextracted rootfs tarballs (GRUB can't untar at boot)
#   2. Refreshes the main grub.cfg (mostly to re-bake the partition UUID
#      in case the drive was re-partitioned)
# Run this after manually copying tarballs, or as a maintenance refresh.
scan_payloads() {
    local payload_mnt="$1"
    local count=0

    # --- Extract unextracted rootfs tarballs ---------------------------------
    # GRUB cannot untar at boot time, so we extract on the host.
    # Once extracted, the resulting directory is auto-discovered by GRUB.
    local tarball name target_dir
    for tarball in "$payload_mnt"/payloads/rootfs_tarballs/*.tar \
                   "$payload_mnt"/payloads/rootfs_tarballs/*.tar.gz \
                   "$payload_mnt"/payloads/rootfs_tarballs/*.tar.xz \
                   "$payload_mnt"/payloads/rootfs_tarballs/*.tar.zst \
                   "$payload_mnt"/payloads/rootfs_tarballs/*.tgz; do
        [[ -f "$tarball" ]] || continue
        name="$(basename "$tarball")"
        name="${name%.tar}"; name="${name%.tar.gz}"; name="${name%.tar.xz}"
        name="${name%.tar.zst}"; name="${name%.tgz}"
        target_dir="$payload_mnt/payloads/rootfs/$name"

        local marker="$target_dir/.ds_extracted"
        if [[ ! -f "$marker" ]]; then
            log_step "Extracting $name -> payloads/rootfs/$name ..."
            rm -rf "$target_dir"
            mkdir -p "$target_dir"
            tar -xf "$tarball" -C "$target_dir" --numeric-owner --checkpoint=100 --checkpoint-action=dot
            touch "$marker"
            echo
            log_info "  ROOT: $name (newly extracted)"
        else
            log_dim "  ROOT: $name (already extracted)"
        fi
        count=$((count + 1))
    done

    # --- Refresh the main grub.cfg (re-bakes UUID) ---------------------------
    generate_grub_main_cfg "$payload_mnt"

    # --- Summarize what's on the drive ---------------------------------------
    # Table-driven count via PAYLOAD_CATEGORIES (single source of truth).
    local row subdir kind label flags
    local -A count_of=()
    for row in "${PAYLOAD_CATEGORIES[@]}"; do
        IFS=':' read -r subdir kind label flags <<<"$row"
        local base="$payload_mnt/payloads/$subdir"
        local n=0
        case "$kind" in
            file)
                # shellcheck disable=SC2206
                local -a globs=($flags)
                local g f
                for g in "${globs[@]}"; do
                    for f in "$base"/$g; do
                        [[ -f "$f" ]] && n=$((n + 1))
                    done
                done
                ;;
            dir)
                local d
                for d in "$base"/*/; do
                    [[ -d "$d" ]] && n=$((n + 1))
                done
                ;;
        esac
        count_of["$label"]=$n
    done

    log_info "Scan complete. GRUB will auto-discover at boot:"
    log_info "  ISOs    : ${count_of[ISOs]}  (in payloads/isos/)"
    log_info "  RootFS  : ${count_of[Extracted RootFS]}  (in payloads/rootfs/)"
    log_info "  Images  : ${count_of[Raw Images]}  (in payloads/images/ — experimental)"
    log_info "  Tarballs: $count  (in payloads/rootfs_tarballs/ — extracted on host)"
    log_dim "  No further action needed. Just boot the USB."
}

# ----------------------------------------------------------------------------
# Payload deployment (copy a file from host -> USB)
# ----------------------------------------------------------------------------

deploy_payload() {
    local disk="$1" src="$2"
    local payload_mnt dest_dir dest_path

    src="$(readlink -f "$src")"
    [[ -f "$src" ]] || die "Source file not found: $src"

    payload_mnt="$(mount_payload_partition "$disk")"
    mount_esp_partition "$disk" "$payload_mnt" >/dev/null

    local bn lower
    bn="$(basename "$src")"
    lower="$(printf '%s' "$bn" | tr '[:upper:]' '[:lower:]')"

    case "$lower" in
        *.iso)
            dest_dir="$payload_mnt/payloads/isos"
            ;;
        *.tar|*.tar.gz|*.tgz|*.tar.xz|*.tar.zst|*.tar.bz2)
            dest_dir="$payload_mnt/payloads/rootfs_tarballs"
            ;;
        *.img|*.raw)
            dest_dir="$payload_mnt/payloads/images"
            ;;
        *)
            log_warn "Unknown file type for '$bn'. Defaulting to /payloads/."
            dest_dir="$payload_mnt/payloads"
            ;;
    esac

    mkdir -p "$dest_dir"
    dest_path="${dest_dir}/${bn}"

    # Free-space check
    local free_kb src_bytes
    free_kb="$(df -P "$dest_dir" 2>/dev/null | awk 'NR==2{print $4}')"
    src_bytes="$(stat -c %s "$src" 2>/dev/null || echo 0)"
    if [[ -n "$free_kb" && "$src_bytes" -gt 0 ]]; then
        local free_bytes=$((free_kb * 1024))
        if (( src_bytes > free_bytes )); then
            die "Insufficient free space on payload partition: need $((src_bytes / 1024 / 1024)) MB, have $((free_bytes / 1024 / 1024)) MB"
        fi
    fi

    log_step "Copying $bn -> ${dest_path#"$payload_mnt"}..."
    # rsync gives progress and resume capability
    rsync -ah --progress "$src" "$dest_path"
    sync

    # For tarballs, extract immediately so GRUB can auto-discover the
    # resulting directory at boot time. ISOs and images need no post-copy
    # step — the boot-time auto-scan will pick them up.
    case "$lower" in
        *.tar|*.tar.gz|*.tgz|*.tar.xz|*.tar.zst|*.tar.bz2)
            local name="$bn"
            name="${name%.tar}"; name="${name%.tar.gz}"; name="${name%.tar.xz}"
            name="${name%.tar.zst}"; name="${name%.tgz}"; name="${name%.tar.bz2}"
            local target_dir="$payload_mnt/payloads/rootfs/$name"
            local marker="$target_dir/.ds_extracted"
            if [[ ! -f "$marker" ]]; then
                log_step "Extracting tarball -> payloads/rootfs/$name ..."
                rm -rf "$target_dir"
                mkdir -p "$target_dir"
                tar -xf "$dest_path" -C "$target_dir" --numeric-owner --checkpoint=100 --checkpoint-action=dot
                touch "$marker"
                echo
            fi
            ;;
    esac

    log_info "Done. '$bn' deployed."
    log_dim "  No scan needed — GRUB auto-discovers payloads at boot time."
    log_dim "  Just eject the drive and boot from it."
}

# ----------------------------------------------------------------------------
# Listing
# ----------------------------------------------------------------------------

# Payload category table — single source of truth for both list_payloads
# and scan_payloads summary. Each row is: "<subdir>:<kind>:<label>:<flags>"
#   <subdir>  path under $payload_mnt/payloads/
#   <kind>    'file' or 'dir'
#   <label>   human-readable heading
#   <flags>   glob extensions only used for 'file' kinds (e.g. '*.iso')
PAYLOAD_CATEGORIES=(
    "isos:file:ISOs:*.iso"
    "rootfs_tarballs:file:RootFS Tarballs:*"
    "rootfs:dir:Extracted RootFS:"
    "images:file:Raw Images:*.img *.raw"
)

list_payloads() {
    local payload_mnt="$1"
    log_step "Registered payloads on this drive:"
    local found=0

    local row subdir kind label flags
    for row in "${PAYLOAD_CATEGORIES[@]}"; do
        IFS=':' read -r subdir kind label flags <<<"$row"
        local base="$payload_mnt/payloads/$subdir"
        local -a matches=()

        # Step-down discovery: build the candidate list, skip if empty.
        case "$kind" in
            file)
                # shellcheck disable=SC2206  # intentional word-splitting on glob list
                local -a globs=($flags)
                local g f
                for g in "${globs[@]}"; do
                    for f in "$base"/$g; do
                        [[ -f "$f" ]] && matches+=("$f")
                    done
                done
                ;;
            dir)
                local d
                for d in "$base"/*/; do
                    [[ -d "$d" ]] && matches+=("${d%/}")
                done
                ;;
        esac
        (( ${#matches[@]} )) || continue

        printf '\n  %s%s:%s\n' "$C_BOLD" "$label" "$C_NC"
        local m size
        for m in "${matches[@]}"; do
            if [[ "$kind" == dir ]]; then
                size="$(du -sh "$m" | cut -f1)"
            else
                size="$(du -h "$m" | cut -f1)"
            fi
            printf '    %-50s %s\n' "$(basename "$m")" "$size"
            ((found++))
        done
    done

    (( found == 0 )) && printf '  %s(empty)%s\n' "$C_DIM" "$C_NC"
    echo
}

# ----------------------------------------------------------------------------
# Cleanup / wipe
# ----------------------------------------------------------------------------

clean_drive() {
    local disk="$1"
    if [[ "${ASSUME_YES:-0}" != "1" ]]; then
        confirm_destructive "$disk" "wipe"
    fi
    unmount_all "$disk" 2>/dev/null || true
    wipe_signatures "$disk"
    log_info "Drive $disk is now blank. Run 'setup' to provision it."
}

# ----------------------------------------------------------------------------
# Helper shell
# ----------------------------------------------------------------------------

open_shell() {
    local disk="$1"
    local payload_mnt esp_mnt
    payload_mnt="$(mount_payload_partition "$disk")"
    esp_mnt="$(mount_esp_partition "$disk" "$payload_mnt")"
    cat <<EOF

${C_BOLD}DriveStage helper shell${C_NC}
  Disk          : $disk
  Payload mount : $payload_mnt
  ESP mount     : $esp_mnt

Useful commands:
  ls \$PAYLOAD_MNT/payloads/
  sudo ./blkstage.sh scan $disk
  sudo ./blkstage.sh list  $disk

Type 'exit' to leave. Partitions will NOT be auto-unmounted.

EOF
    PAYLOAD_MNT="$payload_mnt" ESP_MNT="$esp_mnt" \
        PS1='drivestage> ' bash --noprofile --norc
}

# ----------------------------------------------------------------------------
# Usage / main
# ----------------------------------------------------------------------------

usage() {
    cat <<EOF
${C_BOLD}blkstage.sh${C_NC} — DriveStage Engine (BlkStage)

${C_BOLD}USAGE:${C_NC}
    sudo $0 <command> <device> [options]
    $0 -V, --version    Print version and exit

${C_BOLD}COMMANDS:${C_NC}
    setup  <device> [payload-file]
        Partition, format, install GRUB (UEFI+BIOS), create directory
        structure, write the auto-scan grub.cfg. If a payload file is
        given, deploy it after setup. After this, just drop files into
        payloads/ and boot — no further commands needed.

    deploy <device> <payload-file>
        Copy a .iso / .tar(.gz/.xz/.zst) / .img to the appropriate
        payloads/ subdirectory. Tarballs are extracted on copy.
        ISOs/images need no post-copy step — GRUB auto-discovers at boot.

    scan   <device>
        Maintenance refresh: extract any unextracted tarballs and
        regenerate the auto-scan grub.cfg (re-bakes the partition UUID).
        NOT required after dropping ISOs — only run this if you added
        tarballs manually or repartitioned the drive.

    list   <device>
        Show currently registered payloads and their sizes.

    shell  <device>
        Mount the payload + ESP partitions and drop into a helper shell.

    clean  <device>
        DANGER: wipe the partition table and all signatures on <device>.

${C_BOLD}ENVIRONMENT (override as needed):${C_NC}
    ESP_SIZE_MB       ESP partition size in MiB            (default: 512)
    BIOS_BOOT_SIZE_MB BIOS boot partition size in MiB      (default: 1)
    PAYLOAD_LABEL     ext4 filesystem label                (default: DRIVESTAGE)
    ESP_LABEL         FAT32 ESP label                      (default: DRIVESTAGE-EFI)
    MOUNT_BASE        Mount prefix                         (default: /mnt/drivestage)
    CONF_FILE         Optional config file to source       (default: /etc/drivestage.conf)
    ASSUME_YES        Set to 1 to skip confirmation prompts

${C_BOLD}EXAMPLES:${C_NC}
    # Provision a fresh USB drive at /dev/sdb (one-time)
    sudo $0 setup /dev/sdb

    # Provision + deploy an initial ISO in one go
    sudo $0 setup /dev/sdb ./ubuntu-24.04.iso

    # Add more ISOs later — TWO ways:
    #   (a) Just drag-and-drop .iso files into /mnt/drivestage/payload/payloads/isos/
    #       (or wherever the drive is mounted). No script needed. Boot and go.
    #   (b) Or use deploy for progress bar + tarball auto-extraction:
    sudo $0 deploy /dev/sdb ./archlinux-2026.08.01-x86_64.iso

    # List what's on the drive
    sudo $0 list /dev/sdb

    # Maintenance: extract tarballs + refresh grub.cfg UUID (rarely needed)
    sudo $0 scan /dev/sdb

${C_BOLD}WORKFLOW (drop-and-boot, inspired by Easy2Boot):${C_NC}
    1. sudo $0 setup /dev/sdb              # one-time provisioning
    2. Copy .iso files into payloads/isos/ # drag-and-drop, no script
    3. Eject the USB and boot from it       # GRUB auto-discovers payloads
    4. Pick an entry from the menu          # done!

${C_BOLD}NOTES:${C_NC}
    - The script refuses to operate on the disk hosting / or /boot.
    - Destructive operations require interactive confirmation (or ASSUME_YES=1).
    - BOOT-TIME AUTO-SCAN: GRUB scans payloads/{isos,rootfs,images}/ at boot
      and generates menu entries on the fly. No host-side menu generation.
    - Distros recognized at boot via filename pattern matching:
      Ubuntu, Debian, Arch, Manjaro, Fedora, CentOS Stream, Rocky, Alma,
      openSUSE, Kali, Parrot, Tails, Knoppix, SystemRescue, Alpine, Void,
      Gentoo, Slax, TinyCore, Clonezilla, GParted, Mint, elementary, Pop!_OS.
      Unrecognized ISOs fall back to GRUB's loopback.cfg chainload.
    - Tarballs (.tar/.tar.gz/.tar.xz/.tar.zst) require one-time host-side
      extraction (GRUB can't untar at boot). Use 'deploy' or 'scan'.

EOF
}

# Globals set during setup
CURRENT_DISK=""

main() {
    local cmd="${1:-}" dev="${2:-}"
    [[ -n "$cmd" ]] || { usage; exit 1; }

    case "$cmd" in
        -h|--help|help) usage; exit 0 ;;
        -V|--version) echo "drivestage 0.4.0"; exit 0 ;;
    esac

    check_root "$cmd"
    check_commands

    # Prevent concurrent invocations from racing on mount/unmount/partition operations
    exec 9>/var/lock/drivestage.lock
    if ! flock -n 9; then
        die "Another drivestage instance is running. Wait for it to finish or remove /var/lock/drivestage.lock"
    fi

    # 'setup' may take an optional payload file as $3
    local payload_arg="${3:-}"

    # Resolve device for all commands except 'help'
    if [[ -z "$dev" && "$cmd" != "help" ]]; then
        dev="$(pick_disk_interactive)"
    fi
    local disk=""
    if [[ -n "$dev" ]]; then
        disk="$(resolve_device "$dev")"
        is_whole_disk "$disk" || die "$disk is a partition, not a whole disk."
        guard_root_disk "$disk"
    fi
    CURRENT_DISK="$disk"

    # Cleanup trap: unmount our mountpoints on exit/interrupt.
    # shellcheck disable=SC2154  # rc is assigned inside the trap eval.
    trap 'rc=$?; [[ -n "$CURRENT_DISK" ]] && unmount_all "$CURRENT_DISK" 2>/dev/null || true; exit $rc' EXIT INT TERM

    case "$cmd" in
        setup)
            if [[ "${ASSUME_YES:-0}" != "1" ]]; then
                confirm_destructive "$disk" "partition and format"
            fi
            # Unmount any stale mounts from a previous run before wiping.
            unmount_all "$disk" 2>/dev/null || true
            wipe_signatures "$disk"
            create_partitions "$disk"
            format_partitions "$disk"
            local mounts payload_mnt esp_mnt
            mounts="$(mount_all "$disk")"
            payload_mnt="$(printf '%s' "$mounts" | sed -n '1p')"
            esp_mnt="$(printf '%s' "$mounts" | sed -n '2p')"
            create_directory_structure "$payload_mnt"
            install_grub "$disk" "$payload_mnt" "$esp_mnt"
            scan_payloads "$payload_mnt"
            list_payloads "$payload_mnt"
            if [[ -n "$payload_arg" ]]; then
                log_step "Deploying initial payload: $payload_arg"
                deploy_payload "$disk" "$payload_arg"
            fi
            log_info "Setup complete. Eject $disk safely before booting from it."
            log_dim "Tip: to add more payloads later, copy files to ${payload_mnt}/payloads/ and run: $0 scan $disk"
            ;;
        deploy)
            [[ -n "$payload_arg" ]] || die "deploy requires a payload file argument."
            deploy_payload "$disk" "$payload_arg"
            ;;
        scan)
            local payload_mnt
            payload_mnt="$(mount_payload_partition "$disk")"
            scan_payloads "$payload_mnt"
            list_payloads "$payload_mnt"
            ;;
        list)
            local payload_mnt
            payload_mnt="$(mount_payload_partition "$disk")"
            list_payloads "$payload_mnt"
            ;;
        shell)
            open_shell "$disk"
            ;;
        clean)
            clean_drive "$disk"
            ;;
        *)
            die "Unknown command: $cmd (try: $0 --help)"
            ;;
    esac
}

main "$@"