#!/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 [ $# -lt 2 ]; then
    echo "Usage: sudo $0 /dev/sdX /path/to/ubuntu-24.04.iso [/path/to/winpe-images]"
    echo ""
    echo "  The optional third argument is the path to WinPE deployment content"
    echo "  (e.g., the mounted Media Creator LITE USB). If provided, the images"
    echo "  will be bundled onto the CIDATA partition for automatic import."
    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"
WINPE_SOURCE="${3:-}"

# 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)..."

# Patch ISO to add 'autoinstall' kernel parameter (skips confirmation prompt)
PATCHED_ISO=$(mktemp /tmp/ubuntu-autoinstall-XXXXXX.iso)
cp "$ISO_PATH" "$PATCHED_ISO"
PATCHED_GRUB=$(mktemp /tmp/grub-XXXXXX.cfg)
xorriso -osirrox on -indev "$ISO_PATH" -extract /boot/grub/grub.cfg "$PATCHED_GRUB" 2>/dev/null
sed -i 's|linux\t/casper/vmlinuz  ---|linux\t/casper/vmlinuz autoinstall ---|' "$PATCHED_GRUB"
sed -i 's/^set timeout=30/set timeout=5/' "$PATCHED_GRUB"
xorriso -indev "$PATCHED_ISO" -outdev "$PATCHED_ISO" \
    -map "$PATCHED_GRUB" /boot/grub/grub.cfg \
    -boot_image any replay 2>/dev/null
echo "    Patched GRUB: added 'autoinstall' kernel param, reduced timeout to 5s"

ISO_SIZE=$(stat -c%s "$PATCHED_ISO")
dd if="$PATCHED_ISO" of="$USB_DEV" bs=4M status=progress oflag=sync
sync
rm -f "$PATCHED_ISO" "$PATCHED_GRUB"

# --- 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
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/"

# Copy webapp
WEBAPP_DIR="$SCRIPT_DIR/webapp"
if [ -d "$WEBAPP_DIR" ]; then
    mkdir -p "$MOUNT_POINT/webapp"
    cp -r "$WEBAPP_DIR/app.py" "$WEBAPP_DIR/requirements.txt" "$MOUNT_POINT/webapp/"
    cp -r "$WEBAPP_DIR/templates" "$WEBAPP_DIR/static" "$MOUNT_POINT/webapp/"
    echo "    Copied webapp/"
fi

# Copy pip wheels for offline Flask install
PIP_WHEELS_DIR="$SCRIPT_DIR/pip-wheels"
if [ ! -d "$PIP_WHEELS_DIR" ]; then
    echo "    pip-wheels/ not found — downloading now..."
    mkdir -p "$PIP_WHEELS_DIR"
    if pip3 download -d "$PIP_WHEELS_DIR" flask lxml 2>/dev/null; then
        echo "    Downloaded pip wheels successfully."
    else
        echo "    WARNING: Failed to download pip wheels (no internet?)"
        echo "    The PXE server will need internet to install Flask later,"
        echo "    or manually copy wheels to pip-wheels/ and rebuild."
        rmdir "$PIP_WHEELS_DIR" 2>/dev/null || true
    fi
fi
if [ -d "$PIP_WHEELS_DIR" ]; then
    cp -r "$PIP_WHEELS_DIR" "$MOUNT_POINT/pip-wheels"
    WHEEL_COUNT=$(find "$PIP_WHEELS_DIR" -name '*.whl' | wc -l)
    echo "    Copied pip-wheels/ ($WHEEL_COUNT wheels)"
fi

# Copy WinPE boot files (wimboot, boot.wim, BCD, ipxe.efi, etc.)
BOOT_FILES_DIR="$SCRIPT_DIR/boot-files"
if [ -d "$BOOT_FILES_DIR" ]; then
    BOOT_FILE_COUNT=0
    for bf in "$BOOT_FILES_DIR"/*; do
        if [ -f "$bf" ]; then
            cp "$bf" "$MOUNT_POINT/"
            BOOT_FILE_COUNT=$((BOOT_FILE_COUNT + 1))
        fi
    done
    BOOT_FILES_SIZE=$(du -sh "$BOOT_FILES_DIR" | cut -f1)
    echo "    Copied $BOOT_FILE_COUNT boot files ($BOOT_FILES_SIZE) — wimboot, boot.wim, ipxe.efi, etc."
else
    echo "    WARNING: No boot-files/ found (copy WinPE boot files from Media Creator)"
fi

# Copy boot tools (Clonezilla, Blancco, Memtest) if prepared
BOOT_TOOLS_DIR="$SCRIPT_DIR/boot-tools"
if [ -d "$BOOT_TOOLS_DIR" ]; then
    cp -r "$BOOT_TOOLS_DIR" "$MOUNT_POINT/boot-tools"
    TOOLS_SIZE=$(du -sh "$MOUNT_POINT/boot-tools" | cut -f1)
    echo "    Copied boot-tools/ ($TOOLS_SIZE)"
else
    echo "    No boot-tools/ found (run prepare-boot-tools.sh first)"
fi

# Optionally copy WinPE deployment images
if [ -n "$WINPE_SOURCE" ] && [ -d "$WINPE_SOURCE" ]; then
    echo "    Copying WinPE deployment content from $WINPE_SOURCE..."
    mkdir -p "$MOUNT_POINT/images"
    cp -r "$WINPE_SOURCE"/* "$MOUNT_POINT/images/" 2>/dev/null || true
    IMG_SIZE=$(du -sh "$MOUNT_POINT/images" | cut -f1)
    echo "    Copied WinPE images ($IMG_SIZE)"
elif [ -n "$WINPE_SOURCE" ]; then
    echo "    WARNING: WinPE source path not found: $WINPE_SOURCE (skipping)"
fi

# 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 ""