#!/bin/bash
# bee-install — install the live system to a local disk.
#
# Usage: bee-install <device> [logfile]
#   device  — target block device, e.g. /dev/sda (will be WIPED)
#   logfile — optional path to write progress log (default: /tmp/bee-install.log)
#
# Layout (UEFI):  GPT, /dev/sdX1=EFI 512MB vfat, /dev/sdX2=root ext4
# Layout (BIOS):  MBR, /dev/sdX1=root ext4
#
# Squashfs source: /run/live/medium/live/filesystem.squashfs

set -euo pipefail

usage() {
    cat >&2 <<'EOF'
Usage: bee-install <device> [logfile]

  Installs the live system to a local disk (WIPES the target).

  device   Target block device, e.g. /dev/sda or /dev/nvme0n1
           Must be a hard disk or NVMe — NOT a CD-ROM (/dev/sr*)
  logfile  Optional path for progress log (default: /tmp/bee-install.log)

Examples:
  bee-install /dev/sda
  bee-install /dev/nvme0n1
  bee-install /dev/sdb /tmp/my-install.log

WARNING: ALL DATA ON <device> WILL BE ERASED.

Layout (UEFI):  GPT — partition 1: EFI 512MB vfat, partition 2: root ext4
Layout (BIOS):  MBR — partition 1: root ext4
EOF
    exit 1
}

DEVICE="${1:-}"
LOGFILE="${2:-/tmp/bee-install.log}"

if [ -z "$DEVICE" ] || [ "$DEVICE" = "--help" ] || [ "$DEVICE" = "-h" ]; then
    usage
fi
if [ ! -b "$DEVICE" ]; then
    echo "ERROR: $DEVICE is not a block device" >&2
    echo "Run 'lsblk' to list available disks." >&2
    exit 1
fi
# Block CD-ROM devices
case "$DEVICE" in
    /dev/sr*|/dev/scd*)
        echo "ERROR: $DEVICE is a CD-ROM/optical device — cannot install to it." >&2
        echo "Run 'lsblk' to find the target disk (e.g. /dev/sda, /dev/nvme0n1)." >&2
        exit 1
        ;;
esac
# Check required tools
for tool in parted mkfs.vfat mkfs.ext4 unsquashfs grub-install update-grub; do
    if ! command -v "$tool" >/dev/null 2>&1; then
        echo "ERROR: required tool not found: $tool" >&2
        exit 1
    fi
done

SQUASHFS="/run/live/medium/live/filesystem.squashfs"
if [ ! -f "$SQUASHFS" ]; then
    echo "ERROR: squashfs not found at $SQUASHFS" >&2
    exit 1
fi

MOUNT_ROOT="/mnt/bee-install-root"

# ------------------------------------------------------------------
log() { echo "[$(date +%H:%M:%S)] $*" | tee -a "$LOGFILE"; }
die() { log "ERROR: $*"; exit 1; }

# ------------------------------------------------------------------
# Detect UEFI
if [ -d /sys/firmware/efi ]; then
    UEFI=1
    log "Boot mode: UEFI"
else
    UEFI=0
    log "Boot mode: BIOS/legacy"
fi

# Determine partition names (nvme uses p-suffix: nvme0n1p1)
if echo "$DEVICE" | grep -qE 'nvme|mmcblk'; then
    PART_PREFIX="${DEVICE}p"
else
    PART_PREFIX="${DEVICE}"
fi

if [ "$UEFI" = "1" ]; then
    PART_EFI="${PART_PREFIX}1"
    PART_ROOT="${PART_PREFIX}2"
else
    PART_ROOT="${PART_PREFIX}1"
fi

# ------------------------------------------------------------------
log "=== BEE DISK INSTALLER ==="
log "Target device : $DEVICE"
log "Root partition: $PART_ROOT"
[ "$UEFI" = "1" ] && log "EFI partition : $PART_EFI"
log "Squashfs      : $SQUASHFS ($(du -sh "$SQUASHFS" | cut -f1))"
log "Log           : $LOGFILE"
log ""

# ------------------------------------------------------------------
log "--- Step 1/7: Unmounting target device ---"
# Unmount any partitions on target device
for part in "${DEVICE}"* ; do
    if [ "$part" = "$DEVICE" ]; then continue; fi
    if mountpoint -q "$part" 2>/dev/null; then
        log "  umount $part"
        umount "$part" || true
    fi
done
# Also unmount our mount point if leftover
umount "${MOUNT_ROOT}" 2>/dev/null || true
umount "${MOUNT_ROOT}/boot/efi" 2>/dev/null || true

# ------------------------------------------------------------------
log "--- Step 2/7: Partitioning $DEVICE ---"
if [ "$UEFI" = "1" ]; then
    parted -s "$DEVICE" mklabel gpt
    parted -s "$DEVICE" mkpart EFI fat32 1MiB 513MiB
    parted -s "$DEVICE" set 1 esp on
    parted -s "$DEVICE" mkpart root ext4 513MiB 100%
else
    parted -s "$DEVICE" mklabel msdos
    parted -s "$DEVICE" mkpart primary ext4 1MiB 100%
    parted -s "$DEVICE" set 1 boot on
fi
# Wait for kernel to see new partitions
sleep 1
partprobe "$DEVICE" 2>/dev/null || true
sleep 1
log "  Partitioning done."

# ------------------------------------------------------------------
log "--- Step 3/7: Formatting ---"
if [ "$UEFI" = "1" ]; then
    mkfs.vfat -F32 -n EFI "$PART_EFI"
    log "  $PART_EFI formatted as vfat (EFI)"
fi
mkfs.ext4 -F -L bee-root "$PART_ROOT"
log "  $PART_ROOT formatted as ext4"

# ------------------------------------------------------------------
log "--- Step 4/7: Mounting target ---"
mkdir -p "$MOUNT_ROOT"
mount "$PART_ROOT" "$MOUNT_ROOT"
if [ "$UEFI" = "1" ]; then
    mkdir -p "${MOUNT_ROOT}/boot/efi"
    mount "$PART_EFI" "${MOUNT_ROOT}/boot/efi"
fi
log "  Mounted."

# ------------------------------------------------------------------
log "--- Step 5/7: Unpacking filesystem (this takes 10-20 minutes) ---"
log "  Source: $SQUASHFS"
log "  Target: $MOUNT_ROOT"
unsquashfs -f -d "$MOUNT_ROOT" "$SQUASHFS" 2>&1 | \
    grep -E '^\[|^inod|^created|^extract' | \
    while read -r line; do log "  $line"; done || true
log "  Unpack complete."

# ------------------------------------------------------------------
log "--- Step 6/7: Configuring installed system ---"

# Write /etc/fstab
ROOT_UUID=$(blkid -s UUID -o value "$PART_ROOT")
log "  Root UUID: $ROOT_UUID"
cat > "${MOUNT_ROOT}/etc/fstab" <<FSTAB
# Generated by bee-install
UUID=${ROOT_UUID}  /        ext4  defaults,errors=remount-ro  0  1
tmpfs              /tmp     tmpfs defaults,size=512m          0  0
FSTAB
if [ "$UEFI" = "1" ]; then
    EFI_UUID=$(blkid -s UUID -o value "$PART_EFI")
    echo "UUID=${EFI_UUID}   /boot/efi  vfat  umask=0077  0  1" >> "${MOUNT_ROOT}/etc/fstab"
fi
log "  fstab written."

# Remove live-boot persistence markers so installed system boots normally
rm -f "${MOUNT_ROOT}/etc/live/boot.conf" 2>/dev/null || true
rm -f "${MOUNT_ROOT}/etc/live/live.conf" 2>/dev/null || true

# Bind mount virtual filesystems for chroot
mount --bind /dev  "${MOUNT_ROOT}/dev"
mount --bind /proc "${MOUNT_ROOT}/proc"
mount --bind /sys  "${MOUNT_ROOT}/sys"
[ "$UEFI" = "1" ] && mount --bind /sys/firmware/efi/efivars "${MOUNT_ROOT}/sys/firmware/efi/efivars" 2>/dev/null || true

# ------------------------------------------------------------------
log "--- Step 7/7: Installing GRUB bootloader ---"

# Helper: run a chroot command, log all output, return its exit code.
# Needed because "cmd | while" pipelines hide the exit code of cmd.
chroot_log() {
    local rc=0
    local out
    out=$(chroot "$MOUNT_ROOT" "$@" 2>&1) || rc=$?
    echo "$out" | while IFS= read -r line; do log "  $line"; done
    return $rc
}

if [ "$UEFI" = "1" ]; then
    # Primary attempt: write EFI NVRAM entry (requires writable efivars)
    if ! chroot_log grub-install \
            --target=x86_64-efi \
            --efi-directory=/boot/efi \
            --bootloader-id=bee \
            --recheck; then
        log "  WARNING: grub-install (with NVRAM) failed — retrying with --no-nvram"
        # --no-nvram: write grubx64.efi but skip EFI variable update.
        # Needed on headless servers where efivars is read-only or unavailable.
        chroot_log grub-install \
            --target=x86_64-efi \
            --efi-directory=/boot/efi \
            --bootloader-id=bee \
            --no-nvram \
            --recheck || log "  WARNING: grub-install --no-nvram also failed — check logs"
    fi

    # Always install the UEFI fallback path EFI/BOOT/BOOTX64.EFI.
    # Many UEFI implementations (especially server BMCs and some firmware)
    # ignore the NVRAM boot entry and only look for this path.
    GRUB_EFI="${MOUNT_ROOT}/boot/efi/EFI/bee/grubx64.efi"
    FALLBACK_DIR="${MOUNT_ROOT}/boot/efi/EFI/BOOT"
    if [ -f "$GRUB_EFI" ]; then
        mkdir -p "$FALLBACK_DIR"
        cp "$GRUB_EFI" "${FALLBACK_DIR}/BOOTX64.EFI"
        log "  Fallback EFI binary installed: EFI/BOOT/BOOTX64.EFI"
    else
        log "  WARNING: grubx64.efi not found at $GRUB_EFI — UEFI fallback path not set"
    fi
else
    chroot_log grub-install \
        --target=i386-pc \
        --recheck \
        "$DEVICE" || log "  WARNING: grub-install (BIOS) failed — check logs"
fi

chroot_log update-grub || log "  WARNING: update-grub failed — check logs"
log "  GRUB step complete."

# ------------------------------------------------------------------
# Cleanup
log "--- Cleanup ---"
umount "${MOUNT_ROOT}/sys/firmware/efi/efivars" 2>/dev/null || true
umount "${MOUNT_ROOT}/sys"  2>/dev/null || true
umount "${MOUNT_ROOT}/proc" 2>/dev/null || true
umount "${MOUNT_ROOT}/dev"  2>/dev/null || true
[ "$UEFI" = "1" ] && umount "${MOUNT_ROOT}/boot/efi" 2>/dev/null || true
umount "$MOUNT_ROOT" 2>/dev/null || true
rmdir "$MOUNT_ROOT" 2>/dev/null || true

log ""
log "=== INSTALLATION COMPLETE ==="
log "Remove the ISO and reboot to start the installed system."
