#!/bin/bash
#
# build-usb.sh — Build a bootable PXE-server installer USB
#
# Creates a two-partition USB:
#   Partition 1: Ubuntu Server 24.04 installer (ISO contents)
#   Partition 2: CIDATA volume (autoinstall config, .debs, playbook)
#
# The target machine boots from this USB, Ubuntu auto-installs with
# cloud-init (user-data/meta-data from CIDATA), installs offline .debs,
# and on first boot runs the Ansible playbook to configure PXE services.
#
# Usage:
#   sudo ./build-usb.sh /dev/sdX /path/to/ubuntu-24.04-live-server-amd64.iso
#
# WARNING: This will ERASE the target USB device.

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
AUTOINSTALL_DIR="$SCRIPT_DIR/autoinstall"
PLAYBOOK_DIR="$SCRIPT_DIR/playbook"
OFFLINE_PKG_DIR="$SCRIPT_DIR/offline-packages"

# --- Validate arguments ---
if [ $# -ne 2 ]; then
    echo "Usage: sudo $0 /dev/sdX /path/to/ubuntu-24.04.iso"
    echo ""
    echo "Available removable devices:"
    lsblk -d -o NAME,SIZE,TRAN,RM | grep -E '^\S+\s+\S+\s+(usb)\s+1'
    exit 1
fi

USB_DEV="$1"
ISO_PATH="$2"

# Safety checks
if [ "$(id -u)" -ne 0 ]; then
    echo "ERROR: Must run as root (sudo)."
    exit 1
fi

if [ ! -b "$USB_DEV" ]; then
    echo "ERROR: $USB_DEV is not a block device."
    exit 1
fi

if [ ! -f "$ISO_PATH" ]; then
    echo "ERROR: ISO not found at $ISO_PATH"
    exit 1
fi

# Verify it's a removable device (safety against wiping system disks)
REMOVABLE=$(lsblk -nd -o RM "$USB_DEV" 2>/dev/null || echo "0")
if [ "$REMOVABLE" != "1" ]; then
    echo "WARNING: $USB_DEV does not appear to be a removable device."
    read -rp "Are you SURE you want to erase $USB_DEV? (type YES): " CONFIRM
    if [ "$CONFIRM" != "YES" ]; then
        echo "Aborted."
        exit 1
    fi
fi

# Verify required source files exist
if [ ! -f "$AUTOINSTALL_DIR/user-data" ]; then
    echo "ERROR: user-data not found at $AUTOINSTALL_DIR/user-data"
    exit 1
fi

if [ ! -f "$AUTOINSTALL_DIR/meta-data" ]; then
    echo "ERROR: meta-data not found at $AUTOINSTALL_DIR/meta-data"
    exit 1
fi

if [ ! -f "$PLAYBOOK_DIR/pxe_server_setup.yml" ]; then
    echo "ERROR: pxe_server_setup.yml not found at $PLAYBOOK_DIR/"
    exit 1
fi

echo "============================================"
echo "PXE Server USB Builder"
echo "============================================"
echo "USB Device : $USB_DEV"
echo "ISO        : $ISO_PATH"
echo "Source Dir : $SCRIPT_DIR"
echo ""
echo "This will ERASE all data on $USB_DEV."
read -rp "Continue? (y/N): " PROCEED
if [[ ! "$PROCEED" =~ ^[Yy]$ ]]; then
    echo "Aborted."
    exit 1
fi

# --- Unmount any existing partitions ---
echo ""
echo "[1/6] Unmounting existing partitions on $USB_DEV..."
for part in "${USB_DEV}"*; do
    umount "$part" 2>/dev/null || true
done

# --- Write ISO to USB ---
echo "[2/6] Writing Ubuntu ISO to $USB_DEV (this may take several minutes)..."
dd if="$ISO_PATH" of="$USB_DEV" bs=4M status=progress oflag=sync
sync

# --- Find the end of the ISO to create CIDATA partition ---
echo "[3/6] Creating CIDATA partition after ISO data..."

# Get ISO size in bytes and calculate the start sector for the new partition
ISO_SIZE=$(stat -c%s "$ISO_PATH")
SECTOR_SIZE=512
# Start the CIDATA partition 1MB after the ISO ends (alignment)
START_SECTOR=$(( (ISO_SIZE / SECTOR_SIZE) + 2048 ))

# Use sfdisk to append a new partition
echo "    ISO size: $((ISO_SIZE / 1024 / 1024)) MB"
echo "    CIDATA partition starts at sector $START_SECTOR"

# Add a new partition using sfdisk --append
echo "${START_SECTOR},+,L" | sfdisk --append "$USB_DEV" --no-reread 2>/dev/null || true
partprobe "$USB_DEV"
sleep 2

# Determine the new partition name (could be sdX3, sdX4, etc.)
CIDATA_PART=""
for part in "${USB_DEV}"[0-9]*; do
    # Find the partition that starts at or after our start sector
    PART_START=$(sfdisk -d "$USB_DEV" 2>/dev/null | grep "$part" | grep -o 'start=[[:space:]]*[0-9]*' | grep -o '[0-9]*')
    if [ -n "$PART_START" ] && [ "$PART_START" -ge "$START_SECTOR" ]; then
        CIDATA_PART="$part"
        break
    fi
done

# Fallback: use the last partition
if [ -z "$CIDATA_PART" ]; then
    CIDATA_PART=$(lsblk -ln -o NAME "$USB_DEV" | tail -1)
    CIDATA_PART="/dev/$CIDATA_PART"
fi

echo "    CIDATA partition: $CIDATA_PART"

# --- Format CIDATA partition ---
echo "[4/6] Formatting $CIDATA_PART as FAT32 (label: CIDATA)..."
mkfs.vfat -F 32 -n CIDATA "$CIDATA_PART"

# --- Mount and copy files ---
echo "[5/6] Copying autoinstall config, packages, and playbook to CIDATA..."
MOUNT_POINT=$(mktemp -d)
mount "$CIDATA_PART" "$MOUNT_POINT"

# Copy cloud-init files
cp "$AUTOINSTALL_DIR/user-data" "$MOUNT_POINT/"
cp "$AUTOINSTALL_DIR/meta-data" "$MOUNT_POINT/"

# Copy offline .deb packages into packages/ subdirectory
mkdir -p "$MOUNT_POINT/packages"
DEB_COUNT=0

if [ -d "$OFFLINE_PKG_DIR" ]; then
    for deb in "$OFFLINE_PKG_DIR"/*.deb; do
        if [ -f "$deb" ]; then
            cp "$deb" "$MOUNT_POINT/packages/"
            DEB_COUNT=$((DEB_COUNT + 1))
        fi
    done
fi
echo "    Copied $DEB_COUNT .deb packages to packages/"

# Copy playbook directory
cp -r "$PLAYBOOK_DIR" "$MOUNT_POINT/playbook"
echo "    Copied playbook/"

# List what's on CIDATA
echo ""
echo "    CIDATA contents:"
ls -lh "$MOUNT_POINT/" | sed 's/^/      /'

# --- Cleanup ---
echo ""
echo "[6/6] Syncing and unmounting..."
sync
umount "$MOUNT_POINT"
rmdir "$MOUNT_POINT"

echo ""
echo "============================================"
echo "USB build complete!"
echo "============================================"
echo ""
echo "Next steps:"
echo "  1. Insert USB into target machine"
echo "  2. Boot from USB (F12 / boot menu)"
echo "  3. Ubuntu will auto-install and configure the PXE server"
echo "  4. After reboot, move the NIC to the isolated PXE network"
echo ""