#!/bin/bash
# rpi-partition-bootware
# Selects a non-MMC block device via dialog, wipes it, and creates
# a single FAT32 EFI System Partition for Raspberry Pi bootware.
#
# Requirements: dialog, parted, lsblk, mkfs.vfat (dosfstools)
#
# Author: Vibe coded with Claude <mozes@slackware>

set -euo pipefail

# ── Colour helpers (for non-dialog output) ──────────────────────────────────
RED='\033[0;31m'
YEL='\033[1;33m'
GRN='\033[0;32m'
NC='\033[0m'

err()  { echo -e "${RED}ERROR:${NC} $*" >&2; }
warn() { echo -e "${YEL}WARN:${NC}  $*"; }
info() { echo -e "${GRN}INFO:${NC}  $*"; }

# ── Dependency check ─────────────────────────────────────────────────────────
for cmd in dialog parted lsblk mkfs.vfat; do
    if ! command -v "$cmd" &>/dev/null; then
        err "Required tool not found: $cmd"
        exit 1
    fi
done

# ── Must run as root ─────────────────────────────────────────────────────────
if [[ $EUID -ne 0 ]]; then
    err "This script must be run as root."
    exit 1
fi

# ── Gather non-MMC block devices ─────────────────────────────────────────────
# lsblk: name, size, model — exclude loop, rom, and mmc* devices
mapfile -t DEVICES < <(
    lsblk -d -o NAME,SIZE,MODEL --noheadings \
    | awk '$1 !~ /^(loop|sr|mmc|zram)/' \
    | awk 'NF'
)

if [[ ${#DEVICES[@]} -eq 0 ]]; then
    dialog --title "No Devices Found" \
           --msgbox "\nNo suitable block devices found.\n\nEnsure a USB or SATA drive is connected." \
           8 50
    clear
    err "No non-MMC block devices available."
    exit 1
fi

# ── Build dialog menu list ────────────────────────────────────────────────────
# Format: tag item tag item ...  (dialog radiolist needs: tag item status)
MENU_ITEMS=()
for entry in "${DEVICES[@]}"; do
    # entry e.g.: "sda    16G    SanDisk_Ultra"
    NAME=$(awk '{print $1}' <<< "$entry")
    SIZE=$(awk '{print $2}' <<< "$entry")
    MODEL=$(awk '{$1=$2=""; sub(/^[[:space:]]+/,"",$0); print}' <<< "$entry")
    [[ -z "$MODEL" ]] && MODEL="Unknown"
    MENU_ITEMS+=("/dev/$NAME" "$SIZE  $MODEL" "off")
done

# ── Present radiolist ─────────────────────────────────────────────────────────
DISK=$(dialog --stdout \
              --title "Raspberry Pi Boot Disk Preparation" \
              --radiolist \
              "\nSelect the target block device.\nThis disk will be COMPLETELY ERASED.\n\nUse arrow keys + Space to select:" \
              18 65 "${#MENU_ITEMS[@]}" \
              "${MENU_ITEMS[@]}")

DIALOG_EXIT=$?

clear

if [[ $DIALOG_EXIT -ne 0 ]] || [[ -z "$DISK" ]]; then
    info "Aborted by user."
    exit 0
fi

# ── Validate selection ────────────────────────────────────────────────────────
if [[ ! -b "$DISK" ]]; then
    err "Selected path is not a block device: $DISK"
    exit 1
fi

# ── Check for existing partitions ─────────────────────────────────────────────
EXISTING_PARTS=$(lsblk -ln -o NAME "$DISK" | tail -n +2)

if [[ -n "$EXISTING_PARTS" ]]; then
    PART_LIST=$(lsblk -ln -o NAME,SIZE,FSTYPE,LABEL "$DISK" | tail -n +2)

    dialog --title "Existing Partitions Detected" \
           --yesno \
"The selected device already contains partitions:\n\n
$PART_LIST\n\n
ALL DATA ON $DISK WILL BE PERMANENTLY DESTROYED.\n\n
Do you wish to continue?" \
           14 65

    CHOICE=$?
    clear

    if [[ $CHOICE -ne 0 ]]; then
        info "Aborted by user. No changes made to $DISK."
        exit 0
    fi
fi

# ── Final confirmation ────────────────────────────────────────────────────────
DISK_INFO=$(lsblk -d -o NAME,SIZE,MODEL --noheadings "$DISK" 2>/dev/null || echo "$DISK")

dialog --title "Final Confirmation" \
       --defaultno \
       --yesno \
"You are about to ERASE:\n\n
  Device : $DISK
  Info   : $DISK_INFO\n\n
A single 2 GiB W95 FAT32 partition (type 0x0b) will be created.\n\n
This CANNOT be undone. Are you absolutely sure?" \
       13 60

FINAL=$?
clear

if [[ $FINAL -ne 0 ]]; then
    info "Aborted by user. No changes made."
    exit 0
fi

# ── Unmount any mounted partitions on the target disk ────────────────────────
info "Unmounting any mounted partitions on $DISK ..."
while IFS= read -r part; do
    if mountpoint -q "/dev/$part" 2>/dev/null || \
       grep -q "^/dev/$part " /proc/mounts 2>/dev/null; then
        warn "Unmounting /dev/$part"
        umount "/dev/$part" 2>/dev/null || true
    fi
done < <(lsblk -ln -o NAME "$DISK" | tail -n +2)

# ── Wipe existing partition table ─────────────────────────────────────────────
echo ""
info "Wiping existing partition table on $DISK ..."
wipefs --all --force "$DISK"
dd if=/dev/zero of="$DISK" bs=512 count=1 status=none         # wipe MBR
sync

# ── Write fresh MBR label ─────────────────────────────────────────────────────
info "Writing new MBR partition table ..."
parted --script "$DISK" mklabel msdos

# ── Create 2 GiB W95 FAT32 partition ─────────────────────────────────────────
echo ""
info "Creating 2 GiB W95 FAT32 partition (type 0x0b) ..."
parted --script "$DISK" mkpart primary fat32 1MiB 2049MiB
parted --script "$DISK" set 1 boot on

# Set partition type byte to 0x0b (W95 FAT32) via sfdisk
sfdisk --part-type "$DISK" 1 b

sync

# ── Inform kernel of new partition layout ────────────────────────────────────
info "Informing kernel of partition changes ..."
partprobe "$DISK" 2>/dev/null || true
sleep 1

# ── Determine partition device node ──────────────────────────────────────────
# Handles both /dev/sda -> /dev/sda1  and  /dev/nvme0n1 -> /dev/nvme0n1p1
if [[ "$DISK" =~ (nvme|mmcblk|loop) ]]; then
    PART="${DISK}p1"
else
    PART="${DISK}1"
fi

# Wait briefly for the node to appear
for i in {1..5}; do
    [[ -b "$PART" ]] && break
    sleep 1
done

if [[ ! -b "$PART" ]]; then
    err "Partition device node $PART did not appear. You may need to run: partprobe $DISK"
    exit 1
fi

# ── Format as FAT32 ───────────────────────────────────────────────────────────
info "Formatting $PART as FAT32 ..."
mkfs.vfat -F 32 "$PART"

# ── Apply volume label ────────────────────────────────────────────────────────
info "Applying volume label: SLKhwm_bw ..."
fatlabel "$PART" "SLKhwm_bw"

sync

# ── Summary ───────────────────────────────────────────────────────────────────
echo ""
echo "─────────────────────────────────────────────────────"
info "Done. Partition layout of $DISK:"
parted "$DISK" print
echo "─────────────────────────────────────────────────────"
echo ""
info "Partition     : $PART"
info "Type          : W95 FAT32 (0x0b)"
info "Filesystem    : FAT32  (label: SLKhwm_bw)"
info "Flags         : boot"
echo ""
info "The Raspberry Pi boot disk is ready."
