#!/bin/bash
# usb-provision.sh — Write KubeSolo OS disk image to USB drive for air-gapped deployments
#
# This tool writes a complete KubeSolo OS disk image to a USB drive,
# creating a bootable device with A/B partitions and data partition.
# Optionally bundles a cloud-init config for first-boot provisioning.
#
# Usage:
#   sudo ./hack/usb-provision.sh <disk-image> <device> [--cloud-init <config.yaml>]
#
# Example:
#   sudo ./hack/usb-provision.sh output/kubesolo-os-0.3.0.img /dev/sdb
#   sudo ./hack/usb-provision.sh output/kubesolo-os-0.3.0.img /dev/sdb --cloud-init my-config.yaml
#
# WARNING: This will DESTROY all data on the target device!
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)"

# Color output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

die() { echo -e "${RED}ERROR: $*${NC}" >&2; exit 1; }
warn() { echo -e "${YELLOW}WARNING: $*${NC}" >&2; }
info() { echo -e "${GREEN}==> $*${NC}" >&2; }

# Parse arguments
IMAGE="${1:?Usage: usb-provision.sh <disk-image> <device> [--cloud-init <config.yaml>]}"
DEVICE="${2:?Usage: usb-provision.sh <disk-image> <device> [--cloud-init <config.yaml>]}"
CLOUD_INIT=""

shift 2
while [ $# -gt 0 ]; do
    case "$1" in
        --cloud-init)
            CLOUD_INIT="${2:?--cloud-init requires a config file path}"
            shift 2
            ;;
        *)
            die "Unknown option: $1"
            ;;
    esac
done

# Validation
[ -f "$IMAGE" ] || die "Disk image not found: $IMAGE"
[ -b "$DEVICE" ] || die "Device not found or not a block device: $DEVICE"
[ -n "$CLOUD_INIT" ] && { [ -f "$CLOUD_INIT" ] || die "Cloud-init config not found: $CLOUD_INIT"; }

# Safety: refuse to write to mounted or system devices
if [ "$(id -u)" -ne 0 ]; then
    die "This script must be run as root (use sudo)"
fi

# Check for mounted partitions on the target device
MOUNTED=$(mount | grep "^${DEVICE}" || true)
if [ -n "$MOUNTED" ]; then
    die "Device ${DEVICE} has mounted partitions. Unmount first:\n${MOUNTED}"
fi

# Refuse to write to common system devices
case "$DEVICE" in
    /dev/sda|/dev/nvme0n1|/dev/vda|/dev/xvda)
        warn "Device $DEVICE looks like a system disk!"
        warn "Make absolutely sure this is the correct USB device."
        ;;
esac

# Show device info
echo "" >&2
info "KubeSolo OS USB Provisioning Tool"
echo "" >&2

# Get device info
DEV_SIZE=""
DEV_MODEL=""
if command -v lsblk >/dev/null 2>&1; then
    DEV_SIZE=$(lsblk -dno SIZE "$DEVICE" 2>/dev/null || echo "unknown")
    DEV_MODEL=$(lsblk -dno MODEL "$DEVICE" 2>/dev/null || echo "unknown")
fi

IMAGE_SIZE=$(du -h "$IMAGE" | cut -f1)
IMAGE_NAME=$(basename "$IMAGE")

echo "  Source image: $IMAGE_NAME ($IMAGE_SIZE)" >&2
echo "  Target device: $DEVICE" >&2
[ -n "$DEV_SIZE" ] && echo "  Device size: $DEV_SIZE" >&2
[ -n "$DEV_MODEL" ] && echo "  Device model: $DEV_MODEL" >&2
[ -n "$CLOUD_INIT" ] && echo "  Cloud-init: $CLOUD_INIT" >&2
echo "" >&2

# Confirmation
echo -e "${RED}WARNING: ALL DATA ON ${DEVICE} WILL BE DESTROYED!${NC}" >&2
echo "" >&2
read -rp "Type 'yes' to continue: " CONFIRM
if [ "$CONFIRM" != "yes" ]; then
    echo "Aborted." >&2
    exit 1
fi

echo "" >&2

# Step 1: Write disk image
info "Writing disk image to ${DEVICE}..."
dd if="$IMAGE" of="$DEVICE" bs=4M status=progress conv=fsync 2>&1

# Step 2: Ensure partition table is re-read
info "Re-reading partition table..."
sync
partprobe "$DEVICE" 2>/dev/null || true
sleep 2

# Step 3: Determine partition naming
# /dev/sdb → /dev/sdb1, /dev/sdb2, etc.
# /dev/nvme0n1 → /dev/nvme0n1p1, /dev/nvme0n1p2, etc.
if [[ "$DEVICE" =~ nvme|loop|mmcblk ]]; then
    PART_PREFIX="${DEVICE}p"
else
    PART_PREFIX="${DEVICE}"
fi

DATA_PART="${PART_PREFIX}4"

# Step 4: Expand data partition to fill remaining space
info "Expanding data partition to fill USB drive..."
if command -v growpart >/dev/null 2>&1; then
    growpart "$DEVICE" 4 2>/dev/null || true
elif command -v parted >/dev/null 2>&1; then
    # Use parted to resize partition 4 to use remaining space
    parted -s "$DEVICE" resizepart 4 100% 2>/dev/null || true
else
    warn "Neither growpart nor parted found — data partition not expanded."
    warn "Install cloud-guest-utils (growpart) or parted to auto-expand."
fi

# Resize the filesystem if the partition was expanded
if [ -b "$DATA_PART" ]; then
    e2fsck -f -y "$DATA_PART" 2>/dev/null || true
    resize2fs "$DATA_PART" 2>/dev/null || true
fi

# Step 5: Inject cloud-init config (optional)
if [ -n "$CLOUD_INIT" ]; then
    info "Injecting cloud-init configuration..."

    MOUNT_DIR=$(mktemp -d /tmp/kubesolo-usb-XXXXXX)
    trap "umount '$MOUNT_DIR' 2>/dev/null || true; rmdir '$MOUNT_DIR' 2>/dev/null || true" EXIT

    if [ -b "$DATA_PART" ]; then
        mount "$DATA_PART" "$MOUNT_DIR"
        mkdir -p "$MOUNT_DIR/etc-kubesolo"
        cp "$CLOUD_INIT" "$MOUNT_DIR/etc-kubesolo/cloud-init.yaml"
        sync
        umount "$MOUNT_DIR"
        info "Cloud-init config written to data partition"
    else
        warn "Data partition $DATA_PART not found — cloud-init not injected"
        warn "You can manually copy the config after first boot"
    fi

    rmdir "$MOUNT_DIR" 2>/dev/null || true
    trap - EXIT
fi

# Step 6: Final sync
sync

echo "" >&2
info "USB provisioning complete!"
echo "" >&2
echo "  Device: $DEVICE" >&2
echo "  Image: $IMAGE_NAME" >&2
[ -n "$CLOUD_INIT" ] && echo "  Cloud-init: injected" >&2
echo "" >&2
echo "  Next steps:" >&2
echo "    1. Remove USB drive safely: sudo eject $DEVICE" >&2
echo "    2. Insert into target device and boot from USB" >&2
echo "    3. KubeSolo OS will start automatically" >&2
[ -n "$CLOUD_INIT" ] && echo "    4. Cloud-init config will apply on first boot" >&2
echo "" >&2