#!/bin/bash
# bee-wipe-disks — erase all physical disks (interactive, confirmation required)
#
# Triggered automatically when the kernel cmdline contains bee.wipe=all.
# Can also be run manually from a root shell.
#
# Wipe strategy:
#   NVMe  — nvme format (ATA-style secure erase, fast)
#   Other — blkdiscard -f (TRIM/UNMAP, fast on SSDs)
#           dd if=/dev/zero  (fallback for HDDs, zeros first+last 32 MB)

set -euo pipefail

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

banner() {
    echo ""
    echo "${RED}╔══════════════════════════════════════════════════════════╗${NC}"
    echo "${RED}║      BEE DISK WIPE — ALL DATA WILL BE DESTROYED          ║${NC}"
    echo "${RED}╚══════════════════════════════════════════════════════════╝${NC}"
    echo ""
}

# ── find boot device to skip ──────────────────────────────────────────────────
live_dev() {
    local src
    src=$(findmnt -n -o SOURCE /run/live/medium 2>/dev/null || true)
    [ -z "$src" ] && return
    # Strip partition suffix: /dev/sdb1 → /dev/sdb, /dev/nvme0n1p1 → /dev/nvme0n1
    echo "$src" | sed 's/p\?[0-9]\+$//'
}

# ── enumerate target disks ────────────────────────────────────────────────────
find_disks() {
    local boot_dev
    boot_dev=$(live_dev)

    lsblk -d -n -o NAME,TYPE,SIZE,MODEL | while read -r name type size model; do
        [ "$type" = "disk" ] || continue
        [ "$size" = "0B"  ] && continue           # empty virtual media

        local dev="/dev/$name"
        [ "$dev" = "$boot_dev" ] && continue      # skip boot device

        printf '%s\t%s\t%s\n' "$dev" "$size" "$model"
    done
}

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

    if echo "$dev" | grep -q '^/dev/nvme'; then
        # NVMe format (ses=1 = user data erase)
        if nvme format --ses=1 "$dev" 2>&1; then
            echo "  ${GRN}nvme format OK${NC}"
            return
        fi
        echo "  nvme format failed, falling back to blkdiscard"
    fi

    if blkdiscard -f "$dev" 2>&1; then
        echo "  ${GRN}blkdiscard OK${NC}"
        return
    fi

    echo "  blkdiscard not supported — zeroing partition tables (HDD fallback)"
    local size_bytes
    size_bytes=$(blockdev --getsize64 "$dev")
    local mb32=$(( 32 * 1024 * 1024 ))

    # Zero first 32 MB (MBR, GPT, filesystem superblocks)
    dd if=/dev/zero of="$dev" bs=4M count=8 conv=fsync status=progress 2>&1 || true

    # Zero last 32 MB (backup GPT)
    if [ "$size_bytes" -gt $(( mb32 * 2 )) ]; then
        local 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

    echo "  ${GRN}done (partition tables zeroed)${NC}"
}

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

mapfile -t DISKS < <(find_disks | awk '{print $1}')

if [ ${#DISKS[@]} -eq 0 ]; then
    echo "No physical disks found (boot device excluded)."
    echo "Nothing to wipe."
    exit 0
fi

echo "Disks to be ${RED}COMPLETELY ERASED${NC}:"
echo ""
find_disks | while IFS=$'\t' read -r dev size model; do
    printf "  ${YEL}%-16s${NC} %8s  %s\n" "$dev" "$size" "$model"
done
echo ""
echo "${RED}WARNING: This is IRREVERSIBLE. All data on the listed disks will be lost.${NC}"
echo ""
printf "Type  YES  to confirm wipe, anything else to abort: "
read -r CONFIRM

if [ "$CONFIRM" != "YES" ]; then
    echo ""
    echo "Aborted — no disks were touched."
    exit 0
fi

echo ""
echo "Starting wipe..."

for dev in "${DISKS[@]}"; do
    wipe_disk "$dev"
done

echo ""
echo "${GRN}=== All disks wiped. ===${NC}"
echo ""
printf "Reboot now to return to the boot menu? [Y/n] "
read -r REBOOT
case "${REBOOT:-Y}" in
    [Nn]*) echo "You can reboot manually when ready." ;;
    *)     echo "Rebooting..."; sleep 2; reboot ;;
esac