#!/bin/sh
# bee-initramfs-wipe — interactive disk wipe running entirely in the initramfs.
# Triggered by bee.wipe=all on the kernel cmdline (via local-premount hook).
# Works before squashfs is mounted, so it runs even when live boot fails.

RED='\033[1;31m'
YEL='\033[1;33m'
GRN='\033[1;32m'
CYN='\033[1;36m'
NC='\033[0m'

p()  { printf '%b\n' "$*"; }
pp() { printf '%b'   "$*"; }

banner() {
    p ""
    p "${RED}╔══════════════════════════════════════════════════════════╗${NC}"
    p "${RED}║            BEE DRIVE WIPE — initramfs stage              ║${NC}"
    p "${RED}╚══════════════════════════════════════════════════════════╝${NC}"
    p ""
}

# ── find boot device ─────────────────────────────────────────────────────────
boot_dev() {
    local label token
    for token in $(cat /proc/cmdline 2>/dev/null); do
        case "$token" in
            live-media-label=*) label="${token#*=}" ;;
        esac
    done
    [ -z "$label" ] && return

    local dev
    dev=$(blkid -L "$label" 2>/dev/null) || return
    # strip partition suffix: /dev/sdb1 → /dev/sdb, /dev/nvme0n1p1 → /dev/nvme0n1
    echo "$dev" | sed 's/p\?[0-9]\+$//'
}

# ── enumerate candidate disks ─────────────────────────────────────────────────
list_disks() {
    local boot
    boot=$(boot_dev)

    lsblk -d -n -o NAME,TYPE,SIZE,MODEL 2>/dev/null | while read -r name type size model; do
        [ "$type" = "disk" ] || continue
        [ "$size" = "0B"  ] && continue
        local dev="/dev/$name"
        [ "$dev" = "$boot" ] && continue
        printf '%s\t%s\t%s\n' "$dev" "$size" "${model:-}"
    done
}

# ── wipe one disk ─────────────────────────────────────────────────────────────
wipe_one() {
    local dev="$1"
    p ""
    p "=== ${YEL}${dev}${NC} ==="

    if echo "$dev" | grep -q '^/dev/nvme'; then
        if nvme format --ses=1 "$dev" 2>&1; then
            p "  ${GRN}nvme format OK${NC}"
            blockdev --flushbufs "$dev" 2>/dev/null || true
            return
        fi
        p "  nvme format failed — falling back to blkdiscard"
    fi

    if blkdiscard -f "$dev" 2>&1; then
        p "  ${GRN}blkdiscard OK${NC}"
        blockdev --flushbufs "$dev" 2>/dev/null || true
        return
    fi

    p "  blkdiscard not supported — zeroing partition tables (HDD fallback)"
    local size_bytes mb32 skip
    size_bytes=$(blockdev --getsize64 "$dev" 2>/dev/null || echo 0)
    mb32=$(( 32 * 1024 * 1024 ))

    dd if=/dev/zero of="$dev" bs=4M count=8 conv=fsync status=progress 2>&1 || true

    if [ "$size_bytes" -gt $(( mb32 * 2 )) ]; then
        skip=$(( (size_bytes - mb32) / (4 * 1024 * 1024) ))
        dd if=/dev/zero of="$dev" bs=4M count=8 seek="$skip" conv=fsync status=progress 2>&1 || true
    fi

    blockdev --flushbufs "$dev" 2>/dev/null || true
    p "  ${GRN}done (partition tables zeroed)${NC}"
}

# ── main ──────────────────────────────────────────────────────────────────────
banner

BOOT=$(boot_dev)
[ -n "$BOOT" ] && p "Boot device (excluded): ${CYN}${BOOT}${NC}\n"

# build indexed list
i=0
DEVS=""
IFS='
'
for line in $(list_disks); do
    i=$(( i + 1 ))
    dev=$(echo "$line" | cut -f1)
    size=$(echo "$line" | cut -f2)
    model=$(echo "$line" | cut -f3)
    DEVS="${DEVS}${i}:${dev}:${size}:${model}
"
    printf "  ${CYN}[%d]${NC}  %-16s %8s  %s\n" "$i" "$dev" "$size" "$model"
done
IFS='
'

if [ "$i" -eq 0 ]; then
    p "\nNo physical disks found (boot device excluded)."
    p "Dropping to shell — type 'exit' to continue boot."
    exec /bin/sh
fi

p ""
pp "Enter numbers to wipe (space-separated), ${YEL}all${NC} for all, ${YEL}q${NC} to abort: "
read -r SELECTION

case "$SELECTION" in
    q|Q|'') p "\nAborted."; exec /bin/sh ;;
esac

# resolve selection → list of devs
SELECTED=""
if [ "$SELECTION" = "all" ] || [ "$SELECTION" = "ALL" ]; then
    SELECTED=$(echo "$DEVS" | grep -v '^$' | cut -d: -f2 | tr '\n' ' ')
else
    for num in $SELECTION; do
        match=$(echo "$DEVS" | grep "^${num}:" | cut -d: -f2)
        if [ -z "$match" ]; then
            p "${RED}Unknown index: ${num}${NC}"; exec /bin/sh
        fi
        SELECTED="${SELECTED}${match} "
    done
fi

SELECTED=$(echo "$SELECTED" | tr -s ' ' | sed 's/ $//')

p ""
p "Selected for wipe: ${YEL}${SELECTED}${NC}"
p "${RED}WARNING: This is IRREVERSIBLE. All data on the selected disks will be lost.${NC}"
p ""
pp "Type  YES  to confirm, anything else to abort: "
read -r CONFIRM

if [ "$CONFIRM" != "YES" ]; then
    p "\nAborted — no disks were touched."
    exec /bin/sh
fi

p "\nStarting wipe..."
for dev in $SELECTED; do
    wipe_one "$dev"
done

sync
p ""
p "${GRN}=== All selected disks wiped and flushed. ===${NC}"
p ""
pp "Press Enter to reboot..."
read -r _
reboot