#!/bin/bash
#
# build-usb.sh - Build a bootable PXE-server installer USB (or disk image)
#
# Creates a two-partition layout:
#   Partition 1: Ubuntu Server 24.04 installer (ISO contents, autoinstall patched)
#   Partition N: CIDATA volume (autoinstall config, .debs, playbook, webapp...)
#
# 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 [options] /dev/sdX /path/to/ubuntu-24.04-live-server-amd64.iso [winpe-dir]
#   ./build-usb.sh [options] --image out.img SIZE /path/to/ubuntu.iso [winpe-dir]
#
# Image mode writes the same layout into a regular FILE and needs no root:
# the CIDATA filesystem is built with mkfs.vfat on an image file and filled
# with mtools (mcopy), then placed after the ISO with sfdisk on the file.
# SIZE is a byte count with an optional K/M/G suffix (1024-based), or
# "auto" to size the image to fit the payload.
#
# Options:
#   --image FILE SIZE   build a disk image file instead of writing a device
#   --payload-root DIR  take the gitignored payload dirs (offline-packages/,
#                       pip-wheels/, boot-files/, boot-tools/, enrollment/,
#                       bios-staging/, drivers-staging/) from DIR instead of
#                       the repo root (e.g. a main checkout next to a worktree)
#   --skip-enrollment   do not copy enrollment/
#   --skip-bios         do not copy bios-staging/
#   --skip-drivers      do not copy drivers-staging/
#   --skip-boot-tools   do not copy boot-tools/
#   --site-vars FILE    put FILE on CIDATA as playbook/site-vars.yml (the
#                       per-site NIC MACs; see playbook/site-vars.example.yml)
#   --yes               do not ask for confirmation (device mode still asks
#                       YES for a non-removable device)
#
# WARNING: device mode ERASES the target USB device.

set -euo pipefail

REPO_ROOT="$(cd "$(dirname "$0")"/.. && pwd)"
AUTOINSTALL_DIR="$REPO_ROOT/autoinstall"
PLAYBOOK_DIR="$REPO_ROOT/playbook"
WEBAPP_DIR="$REPO_ROOT/webapp"

usage() {
    echo "Usage: sudo $0 [options] /dev/sdX /path/to/ubuntu-24.04.iso [/path/to/winpe-images]"
    echo "       $0 [options] --image out.img SIZE|auto /path/to/ubuntu-24.04.iso [/path/to/winpe-images]"
    echo ""
    echo "  Options: --payload-root DIR, --site-vars FILE, --skip-enrollment, --skip-bios,"
    echo "           --skip-drivers, --skip-boot-tools, --yes"
    echo ""
    echo "  The optional last 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 2>/dev/null | grep -E '^\S+\s+\S+\s+(usb)\s+1' || true
}

# --- Parse arguments ---
IMAGE_MODE=0
IMAGE_FILE=""
IMAGE_SIZE_ARG=""
PAYLOAD_ROOT="$REPO_ROOT"
SKIP_ENROLLMENT=0
SKIP_BIOS=0
SKIP_DRIVERS=0
SKIP_BOOT_TOOLS=0
ASSUME_YES=0
SITE_VARS=""
POSITIONAL=()
while [ $# -gt 0 ]; do
    case "$1" in
        --image)
            [ $# -ge 3 ] || { usage; exit 1; }
            IMAGE_MODE=1; IMAGE_FILE="$2"; IMAGE_SIZE_ARG="$3"; shift 3 ;;
        --payload-root)
            [ $# -ge 2 ] || { usage; exit 1; }
            PAYLOAD_ROOT="$(cd "$2" && pwd)"; shift 2 ;;
        --site-vars)
            [ $# -ge 2 ] || { usage; exit 1; }
            SITE_VARS="$2"; shift 2 ;;
        --skip-enrollment) SKIP_ENROLLMENT=1; shift ;;
        --skip-bios) SKIP_BIOS=1; shift ;;
        --skip-drivers) SKIP_DRIVERS=1; shift ;;
        --skip-boot-tools) SKIP_BOOT_TOOLS=1; shift ;;
        --yes|-y) ASSUME_YES=1; shift ;;
        -h|--help) usage; exit 0 ;;
        --) shift; POSITIONAL+=("$@"); break ;;
        -*) echo "ERROR: unknown option $1"; usage; exit 1 ;;
        *) POSITIONAL+=("$1"); shift ;;
    esac
done

if [ "$IMAGE_MODE" -eq 1 ]; then
    if [ ${#POSITIONAL[@]} -lt 1 ]; then usage; exit 1; fi
    USB_DEV=""
    ISO_PATH="${POSITIONAL[0]}"
    WINPE_SOURCE="${POSITIONAL[1]:-}"
else
    if [ ${#POSITIONAL[@]} -lt 2 ]; then usage; exit 1; fi
    USB_DEV="${POSITIONAL[0]}"
    ISO_PATH="${POSITIONAL[1]}"
    WINPE_SOURCE="${POSITIONAL[2]:-}"
fi

OFFLINE_PKG_DIR="$PAYLOAD_ROOT/offline-packages"
PIP_WHEELS_DIR="$PAYLOAD_ROOT/pip-wheels"
BOOT_FILES_DIR="$PAYLOAD_ROOT/boot-files"
BOOT_TOOLS_DIR="$PAYLOAD_ROOT/boot-tools"
ENROLLMENT_DIR="$PAYLOAD_ROOT/enrollment"
BIOS_DIR="$PAYLOAD_ROOT/bios-staging"
DRIVERS_DIR="$PAYLOAD_ROOT/drivers-staging"
[ "$SKIP_ENROLLMENT" -eq 1 ] && ENROLLMENT_DIR=""
[ "$SKIP_BIOS" -eq 1 ] && BIOS_DIR=""
[ "$SKIP_DRIVERS" -eq 1 ] && DRIVERS_DIR=""
[ "$SKIP_BOOT_TOOLS" -eq 1 ] && BOOT_TOOLS_DIR=""

# FAT32 has a 4GB max file size; files larger than this are split into
# chunks that the playbook reassembles with `cat`.
FAT32_MAX=$((3500 * 1024 * 1024))   # 3500 MiB chunks, safely under 4GiB
SECTOR_SIZE=512

# --- Safety checks ---
if [ ! -f "$ISO_PATH" ]; then
    echo "ERROR: ISO not found at $ISO_PATH"
    exit 1
fi
ISO_PATH="$(cd "$(dirname "$ISO_PATH")" && pwd)/$(basename "$ISO_PATH")"

need_tools=(xorriso 7z sfdisk mkfs.vfat)
if [ "$IMAGE_MODE" -eq 1 ]; then
    need_tools+=(mcopy mmd mdir truncate)
else
    need_tools+=(partprobe)
fi
for t in "${need_tools[@]}"; do
    if ! command -v "$t" >/dev/null 2>&1; then
        echo "ERROR: required tool '$t' not found in PATH"
        exit 1
    fi
done

if [ "$IMAGE_MODE" -eq 1 ]; then
    if [ -e "$IMAGE_FILE" ] && [ ! -f "$IMAGE_FILE" ]; then
        echo "ERROR: $IMAGE_FILE exists and is not a regular file (image mode never writes devices)."
        exit 1
    fi
    IMAGE_FILE="$(cd "$(dirname "$IMAGE_FILE")" && pwd)/$(basename "$IMAGE_FILE")"
else
    if [ "$(id -u)" -ne 0 ]; then
        echo "ERROR: Must run as root (sudo). (Use --image to build a file without root.)"
        exit 1
    fi
    if [ ! -b "$USB_DEV" ]; then
        echo "ERROR: $USB_DEV is not a block device."
        exit 1
    fi
    # Verify it's a removable device (safety against wiping system disks)
    REMOVABLE=$(lsblk -nd -o RM "$USB_DEV" 2>/dev/null | tr -d '[:space:]' || 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
fi

if [ -n "$SITE_VARS" ] && [ ! -f "$SITE_VARS" ]; then
    echo "ERROR: --site-vars file not found: $SITE_VARS"
    exit 1
fi

# Verify required source files exist
for f in "$AUTOINSTALL_DIR/user-data" "$AUTOINSTALL_DIR/meta-data" "$PLAYBOOK_DIR/pxe_server_setup.yml"; do
    if [ ! -f "$f" ]; then
        echo "ERROR: $f not found"
        exit 1
    fi
done

echo "============================================"
echo "PXE Server USB Builder"
echo "============================================"
if [ "$IMAGE_MODE" -eq 1 ]; then
    echo "Image file : $IMAGE_FILE (size: $IMAGE_SIZE_ARG)"
else
    echo "USB Device : $USB_DEV"
fi
echo "ISO        : $ISO_PATH"
echo "Source Dir : $REPO_ROOT"
echo "Payloads   : $PAYLOAD_ROOT"
echo ""
if [ "$ASSUME_YES" -ne 1 ]; then
    if [ "$IMAGE_MODE" -eq 1 ]; then
        echo "This will overwrite $IMAGE_FILE."
    else
        echo "This will ERASE all data on $USB_DEV."
    fi
    read -rp "Continue? (y/N): " PROCEED
    if [[ ! "$PROCEED" =~ ^[Yy]$ ]]; then
        echo "Aborted."
        exit 1
    fi
fi

# Work dir: next to the output image in image mode (multi-GB, keep it off
# a tmpfs /tmp), else the default temp dir.
if [ "$IMAGE_MODE" -eq 1 ]; then
    WORK=$(mktemp -d "$(dirname "$IMAGE_FILE")/.build-usb.XXXXXX")
else
    WORK=$(mktemp -d)
fi
MOUNT_POINT=""
cleanup() {
    if [ -n "$MOUNT_POINT" ] && mountpoint -q "$MOUNT_POINT" 2>/dev/null; then
        umount "$MOUNT_POINT" || true
    fi
    [ -n "$MOUNT_POINT" ] && rmdir "$MOUNT_POINT" 2>/dev/null || true
    rm -rf "$WORK"
}
trap cleanup EXIT

# --- Unmount any existing partitions ---
echo ""
if [ "$IMAGE_MODE" -eq 0 ]; then
    echo "[1/6] Unmounting existing partitions on $USB_DEV..."
    for part in "${USB_DEV}"*; do
        umount "$part" 2>/dev/null || true
    done
else
    echo "[1/6] Image mode: nothing to unmount"
fi

# --- Rebuild ISO with 'autoinstall' kernel parameter ---
echo "[2/6] Rebuilding ISO with autoinstall kernel parameter..."

# Volume label and El Torito load sizes come from the source ISO, so a
# point release (24.04.3, 24.04.5, ...) does not need a script edit.
ELTORITO_REPORT=$(xorriso -indev "$ISO_PATH" -report_el_torito as_mkisofs 2>/dev/null)
ISO_VOLID=$(printf '%s\n' "$ELTORITO_REPORT" | sed -n "s/^-V '\(.*\)'\$/\1/p" | head -1)
# Load sizes in order: first follows -b (BIOS), last follows -e (EFI).
mapfile -t LOAD_SIZES < <(printf '%s\n' "$ELTORITO_REPORT" | sed -n 's/^-boot-load-size \([0-9][0-9]*\)$/\1/p')
if [ -z "$ISO_VOLID" ] || [ ${#LOAD_SIZES[@]} -lt 2 ]; then
    echo "ERROR: could not read volume id / El Torito load sizes from $ISO_PATH"
    echo "       (xorriso -indev ISO -report_el_torito as_mkisofs)"
    exit 1
fi
BIOS_LOAD_SIZE="${LOAD_SIZES[0]}"
EFI_LOAD_SIZE="${LOAD_SIZES[${#LOAD_SIZES[@]}-1]}"
echo "    Volume id: $ISO_VOLID (boot-load-size BIOS=$BIOS_LOAD_SIZE EFI=$EFI_LOAD_SIZE)"

7z -y x "$ISO_PATH" -o"$WORK/iso" >/dev/null 2>&1
mv "$WORK/iso/[BOOT]" "$WORK/BOOT"
chmod -R u+w "$WORK/iso"

# Patch grub.cfg: add 'autoinstall' to kernel cmdline, reduce timeout
GRUB_CFG="$WORK/iso/boot/grub/grub.cfg"
sed -i 's|linux\t/casper/vmlinuz  ---|linux\t/casper/vmlinuz autoinstall ---|' "$GRUB_CFG"
sed -i 's/^set timeout=30/set timeout=5/' "$GRUB_CFG"
if ! grep -q 'vmlinuz.*autoinstall' "$GRUB_CFG"; then
    echo "ERROR: grub.cfg patch failed: no 'autoinstall' on a kernel line in boot/grub/grub.cfg"
    echo "       The ISO's grub.cfg format changed; update the sed in this script."
    exit 1
fi

PATCHED_ISO="$WORK/patched.iso"
(
    cd "$WORK/iso"
    xorriso -as mkisofs -r \
        -V "$ISO_VOLID" \
        -o "$PATCHED_ISO" \
        --grub2-mbr ../BOOT/1-Boot-NoEmul.img \
        --protective-msdos-label \
        -partition_cyl_align off \
        -partition_offset 16 \
        --mbr-force-bootable \
        -append_partition 2 28732ac11ff8d211ba4b00a0c93ec93b ../BOOT/2-Boot-NoEmul.img \
        -appended_part_as_gpt \
        -iso_mbr_part_type a2a0d0ebe5b9334487c068b6b72699c7 \
        -c '/boot.catalog' \
        -b '/boot/grub/i386-pc/eltorito.img' \
        -no-emul-boot \
        -boot-load-size "$BIOS_LOAD_SIZE" \
        -boot-info-table \
        --grub2-boot-info \
        -eltorito-alt-boot \
        -e '--interval:appended_partition_2:::' \
        -no-emul-boot \
        -boot-load-size "$EFI_LOAD_SIZE" \
        . 2>"$WORK/xorriso.log"
) || { echo "ERROR: xorriso failed:"; tail -20 "$WORK/xorriso.log"; exit 1; }
rm -rf "$WORK/iso"
echo "    ISO rebuilt with 'autoinstall' kernel param and 5s GRUB timeout"

ISO_SIZE=$(stat -c%s "$PATCHED_ISO")
# Start the CIDATA partition 1MB after the ISO ends (alignment)
START_SECTOR=$(( (ISO_SIZE / SECTOR_SIZE) + 2048 ))

# --- Payload size (image mode sizing) ---
src_bytes() {
    # Apparent size in bytes of the given paths (missing ones count 0)
    local total=0 p n
    for p in "$@"; do
        [ -n "$p" ] && [ -e "$p" ] || continue
        n=$(du -sbL "$p" 2>/dev/null | cut -f1)
        total=$((total + ${n:-0}))
    done
    echo "$total"
}

parse_size() {
    local s="$1" num unit
    num="${s%[KkMmGgTt]}"
    unit="${s#"$num"}"
    [[ "$num" =~ ^[0-9]+$ ]] || { echo "ERROR: bad size '$s'" >&2; return 1; }
    case "$unit" in
        "") echo "$num" ;;
        K|k) echo $((num * 1024)) ;;
        M|m) echo $((num * 1024 * 1024)) ;;
        G|g) echo $((num * 1024 * 1024 * 1024)) ;;
        T|t) echo $((num * 1024 * 1024 * 1024 * 1024)) ;;
    esac
}

if [ "$IMAGE_MODE" -eq 1 ]; then
    PAYLOAD_BYTES=$(src_bytes "$AUTOINSTALL_DIR/user-data" "$AUTOINSTALL_DIR/meta-data" \
        "$OFFLINE_PKG_DIR" "$PLAYBOOK_DIR" "$WEBAPP_DIR" "$PIP_WHEELS_DIR" \
        "$BOOT_FILES_DIR" "$BOOT_TOOLS_DIR" "$ENROLLMENT_DIR" "$BIOS_DIR" \
        "$DRIVERS_DIR" "$WINPE_SOURCE")
    if [ "$IMAGE_SIZE_ARG" = "auto" ]; then
        # payload + 10% + 256 MiB FAT/cluster slack, rounded up to 1 MiB
        CIDATA_BYTES=$(( PAYLOAD_BYTES + PAYLOAD_BYTES / 10 + 256 * 1024 * 1024 ))
        IMAGE_BYTES=$(( START_SECTOR * SECTOR_SIZE + CIDATA_BYTES + 2 * 1024 * 1024 ))
    else
        IMAGE_BYTES=$(parse_size "$IMAGE_SIZE_ARG")
    fi
    IMAGE_BYTES=$(( (IMAGE_BYTES + 1048575) / 1048576 * 1048576 ))
    MIN_BYTES=$(( START_SECTOR * SECTOR_SIZE + PAYLOAD_BYTES + 64 * 1024 * 1024 ))
    if [ "$IMAGE_BYTES" -lt "$MIN_BYTES" ]; then
        echo "ERROR: image size $((IMAGE_BYTES / 1048576)) MiB too small; need at least $((MIN_BYTES / 1048576)) MiB"
        echo "       (ISO $((ISO_SIZE / 1048576)) MiB + payload $((PAYLOAD_BYTES / 1048576)) MiB). Use SIZE=auto."
        exit 1
    fi
    echo "    Payload: $((PAYLOAD_BYTES / 1048576)) MiB, image: $((IMAGE_BYTES / 1048576)) MiB"
fi

# --- Write the ISO ---
if [ "$IMAGE_MODE" -eq 1 ]; then
    echo "    Writing patched ISO to $IMAGE_FILE..."
    rm -f "$IMAGE_FILE"
    truncate -s "$IMAGE_BYTES" "$IMAGE_FILE"
    dd if="$PATCHED_ISO" of="$IMAGE_FILE" bs=4M conv=notrunc,sparse status=none
    TARGET="$IMAGE_FILE"
else
    echo "    Writing patched ISO to $USB_DEV..."
    dd if="$PATCHED_ISO" of="$USB_DEV" bs=4M status=progress oflag=sync
    sync
    TARGET="$USB_DEV"
fi
rm -f "$PATCHED_ISO"

# --- Create the CIDATA partition after the ISO data ---
echo "[3/6] Creating CIDATA partition after ISO data..."
echo "    ISO size: $((ISO_SIZE / 1024 / 1024)) MB"
echo "    CIDATA partition starts at sector $START_SECTOR"

# A failed append must stop the build: the old fallback ("last partition")
# then formatted the ISO's own EFI partition.
if ! echo "${START_SECTOR},+,L" | sfdisk --append "$TARGET" --no-reread >"$WORK/sfdisk.log" 2>&1; then
    echo "ERROR: sfdisk --append failed on $TARGET:"
    sed 's/^/      /' "$WORK/sfdisk.log"
    exit 1
fi

# Find the partition that starts exactly at our start sector.
# sfdisk -d lines look like: "/dev/sdb3 : start=  6324224, size= ..., type=..."
read_new_part() {
    sfdisk -d "$TARGET" 2>/dev/null | awk -v want="$START_SECTOR" '
        / : start=/ {
            node = $1
            line = $0
            sub(/.*start=[ \t]*/, "", line); start = line; sub(/[^0-9].*/, "", start)
            line = $0
            sub(/.*size=[ \t]*/, "", line); size = line; sub(/[^0-9].*/, "", size)
            if (start + 0 == want + 0) { print node, start, size; exit }
        }'
}
NEW_PART=$(read_new_part)
if [ -z "$NEW_PART" ]; then
    echo "ERROR: no partition starting at sector $START_SECTOR after sfdisk --append. Table:"
    sfdisk -d "$TARGET" 2>&1 | sed 's/^/      /'
    exit 1
fi
read -r PART_NODE PART_START PART_SECTORS <<<"$NEW_PART"

if [ "$IMAGE_MODE" -eq 0 ]; then
    partprobe "$USB_DEV"
    sleep 2
    CIDATA_PART="$PART_NODE"
    if [ ! -b "$CIDATA_PART" ]; then
        echo "ERROR: $CIDATA_PART (from sfdisk) is not a block device after partprobe"
        exit 1
    fi
    # Refuse to format anything that starts before the computed start sector
    # (the ISO9660 data and its EFI partition live there).
    SYS_START=$(cat "/sys/class/block/$(basename "$CIDATA_PART")/start" 2>/dev/null || echo 0)
    if [ "$SYS_START" -lt "$START_SECTOR" ]; then
        echo "ERROR: refusing to format $CIDATA_PART: it starts at sector $SYS_START, before $START_SECTOR"
        exit 1
    fi
    echo "    CIDATA partition: $CIDATA_PART"
else
    if [ "$PART_START" -lt "$START_SECTOR" ]; then
        echo "ERROR: refusing to format partition at sector $PART_START (< $START_SECTOR)"
        exit 1
    fi
    echo "    CIDATA partition: $PART_NODE (sectors $PART_START + $PART_SECTORS)"
fi

# --- Format CIDATA partition ---
if [ "$IMAGE_MODE" -eq 0 ]; then
    echo "[4/6] Formatting $CIDATA_PART as FAT32 (label: CIDATA)..."
    mkfs.vfat -F 32 -n CIDATA "$CIDATA_PART"
    MOUNT_POINT=$(mktemp -d)
    mount "$CIDATA_PART" "$MOUNT_POINT"
else
    echo "[4/6] Building FAT32 CIDATA filesystem image ($((PART_SECTORS / 2048)) MiB)..."
    CIDATA_IMG="$WORK/cidata.img"
    # mkfs.vfat -C takes the size in 1 KiB blocks
    mkfs.vfat -F 32 -n CIDATA -C "$CIDATA_IMG" $((PART_SECTORS / 2)) >/dev/null
    export MTOOLS_SKIP_CHECK=1
fi

# --- Target-neutral copy helpers ---
# Paths are relative to the CIDATA root ("" = root).
t_mkdir() {
    local rel="$1" acc="" part
    if [ "$IMAGE_MODE" -eq 0 ]; then
        mkdir -p "$MOUNT_POINT/$rel"
        return
    fi
    local IFS=/
    for part in $rel; do
        [ -n "$part" ] || continue
        acc="$acc/$part"
        mmd -i "$CIDATA_IMG" "::$acc" 2>/dev/null || true
    done
}
# t_cp DEST_REL_DIR SRC... : copy files/dirs (recursively) into a CIDATA dir
t_cp() {
    local dest="$1"; shift
    if [ "$IMAGE_MODE" -eq 0 ]; then
        cp -r "$@" "$MOUNT_POINT/$dest/"
    else
        mcopy -i "$CIDATA_IMG" -s -Q -D o "$@" "::/${dest:+$dest/}"
    fi
}
# t_cp_split SRC DEST_REL_DIR : copy one file, splitting it if over FAT32_MAX
SPLIT_COUNT=0
t_cp_split() {
    local f="$1" dest="$2" bn
    bn="$(basename "$f")"
    if [ "$(stat -L -c%s "$f")" -gt "$FAT32_MAX" ]; then
        echo "      Splitting $bn (>$((FAT32_MAX / 1024 / 1024))M) into chunks..."
        if [ "$IMAGE_MODE" -eq 0 ]; then
            split -b "$FAT32_MAX" -d -a 2 "$f" "$MOUNT_POINT/$dest/${bn}.part."
        else
            mkdir -p "$WORK/split"
            split -b "$FAT32_MAX" -d -a 2 "$f" "$WORK/split/${bn}.part."
            t_cp "$dest" "$WORK/split/"*
            rm -rf "$WORK/split"
        fi
        SPLIT_COUNT=$((SPLIT_COUNT + 1))
    else
        t_cp "$dest" "$f"
    fi
}

# --- Copy files ---
echo "[5/6] Copying autoinstall config, packages, and playbook to CIDATA..."

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

# Copy offline .deb packages into packages/ subdirectory
t_mkdir packages
DEB_COUNT=0
if [ -d "$OFFLINE_PKG_DIR" ] && compgen -G "$OFFLINE_PKG_DIR/*.deb" >/dev/null; then
    t_cp packages "$OFFLINE_PKG_DIR"/*.deb
    DEB_COUNT=$(find "$OFFLINE_PKG_DIR" -maxdepth 1 -name '*.deb' | wc -l)
else
    echo "    WARNING: no .deb packages in $OFFLINE_PKG_DIR (run download-packages.sh)"
fi
echo "    Copied $DEB_COUNT .deb packages to packages/"

# Copy playbook directory
t_cp "" "$PLAYBOOK_DIR"
echo "    Copied playbook/"
if [ -n "$SITE_VARS" ]; then
    mkdir -p "$WORK/sv"
    cp "$SITE_VARS" "$WORK/sv/site-vars.yml"
    t_cp playbook "$WORK/sv/site-vars.yml"
    rm -rf "$WORK/sv"
    echo "    Copied $SITE_VARS as playbook/site-vars.yml"
elif [ ! -f "$PLAYBOOK_DIR/site-vars.yml" ]; then
    echo "    WARNING: no playbook/site-vars.yml (and no --site-vars): the playbook"
    echo "             stops until pxe_nic_mac / corp_nic_mac are set on the server."
fi

# Copy the whole webapp tree (app.py imports config, services/, views/;
# the Blancco evidence timer runs blancco_evidence_monthly.py), minus caches.
if [ -d "$WEBAPP_DIR" ]; then
    mkdir -p "$WORK/stage"
    cp -r "$WEBAPP_DIR" "$WORK/stage/webapp"
    find "$WORK/stage/webapp" -name __pycache__ -type d -prune -exec rm -rf {} +
    find "$WORK/stage/webapp" -name '*.py[co]' -delete
    t_cp "" "$WORK/stage/webapp"
    rm -rf "$WORK/stage"
    echo "    Copied webapp/"
fi

# Copy pip wheels for offline Flask install
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
    t_cp "" "$PIP_WHEELS_DIR"
    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.) to the root
if [ -d "$BOOT_FILES_DIR" ]; then
    BOOT_FILE_COUNT=0
    for bf in "$BOOT_FILES_DIR"/*; do
        if [ -f "$bf" ]; then
            t_cp_split "$bf" ""
            BOOT_FILE_COUNT=$((BOOT_FILE_COUNT + 1))
        fi
    done
    BOOT_FILES_SIZE=$(du -shL "$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
if [ -n "$BOOT_TOOLS_DIR" ] && [ -d "$BOOT_TOOLS_DIR" ]; then
    t_cp "" "$BOOT_TOOLS_DIR"
    TOOLS_SIZE=$(du -shL "$BOOT_TOOLS_DIR" | cut -f1)
    echo "    Copied boot-tools/ ($TOOLS_SIZE)"
elif [ -z "$BOOT_TOOLS_DIR" ]; then
    echo "    Skipped boot-tools/ (--skip-boot-tools)"
else
    echo "    No boot-tools/ found (run prepare-boot-tools.sh first)"
fi

# Copy enrollment directory (PPKGs, run-enrollment.ps1) if present.
# Files over the FAT32 limit are split; the playbook reassembles them.
if [ -n "$ENROLLMENT_DIR" ] && [ -d "$ENROLLMENT_DIR" ]; then
    t_mkdir enrollment
    SPLIT_COUNT=0
    for f in "$ENROLLMENT_DIR"/*; do
        [ -e "$f" ] || continue
        if [ -f "$f" ]; then
            t_cp_split "$f" enrollment
        else
            t_cp enrollment "$f"
        fi
    done
    PPKG_COUNT=$(find "$ENROLLMENT_DIR" -maxdepth 1 -name '*.ppkg' 2>/dev/null | wc -l)
    ENROLL_SIZE=$(du -shL "$ENROLLMENT_DIR" | cut -f1)
    echo "    Copied enrollment/ ($ENROLL_SIZE, $PPKG_COUNT PPKGs, $SPLIT_COUNT split)"
elif [ -z "$ENROLLMENT_DIR" ]; then
    echo "    Skipped enrollment/ (--skip-enrollment)"
else
    echo "    No enrollment/ directory found (PPKGs can be uploaded via webapp later)"
fi

# Copy BIOS update binaries if staged
if [ -n "$BIOS_DIR" ] && [ -d "$BIOS_DIR" ] && [ "$(ls -A "$BIOS_DIR" 2>/dev/null)" ]; then
    echo "    Copying BIOS update binaries from bios-staging/..."
    t_mkdir bios
    t_cp bios "$BIOS_DIR"/* || true
    BIOS_COUNT=$(find "$BIOS_DIR" -name '*.exe' 2>/dev/null | wc -l)
    BIOS_SIZE=$(du -shL "$BIOS_DIR" | cut -f1)
    echo "    Copied bios/ ($BIOS_SIZE, $BIOS_COUNT files)"
elif [ -z "$BIOS_DIR" ]; then
    echo "    Skipped bios-staging/ (--skip-bios)"
else
    echo "    No bios-staging/ found (BIOS updates can be pushed via download-drivers.py later)"
fi

# Copy Dell driver packs if staged (split over the FAT32 limit)
if [ -n "$DRIVERS_DIR" ] && [ -d "$DRIVERS_DIR" ] && [ "$(ls -A "$DRIVERS_DIR" 2>/dev/null)" ]; then
    echo "    Copying Dell driver packs from drivers-staging/..."
    t_mkdir drivers
    SPLIT_COUNT=0
    # Mirror directory tree first
    while IFS= read -r d; do
        t_mkdir "drivers/${d#./}"
    done < <(cd "$DRIVERS_DIR" && find . -mindepth 1 -type d)
    while IFS= read -r f; do
        rel="${f#"$DRIVERS_DIR"/}"
        reldir="$(dirname "$rel")"
        [ "$reldir" = "." ] && reldir="" || reldir="/$reldir"
        t_cp_split "$f" "drivers$reldir"
    done < <(find "$DRIVERS_DIR" -type f)
    DRIVERS_SIZE=$(du -shL "$DRIVERS_DIR" | cut -f1)
    echo "    Copied drivers/ ($DRIVERS_SIZE, $SPLIT_COUNT split)"
elif [ -z "$DRIVERS_DIR" ]; then
    echo "    Skipped drivers-staging/ (--skip-drivers)"
else
    echo "    No drivers-staging/ found (drivers can be downloaded later)"
fi

# Optionally copy WinPE deployment images
if [ -n "$WINPE_SOURCE" ] && [ -d "$WINPE_SOURCE" ]; then
    echo "    Copying WinPE deployment content from $WINPE_SOURCE..."
    t_mkdir images
    t_cp images "$WINPE_SOURCE"/* || true
    IMG_SIZE=$(du -shL "$WINPE_SOURCE" | 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:"
if [ "$IMAGE_MODE" -eq 0 ]; then
    # shellcheck disable=SC2012  # human-readable listing only
    ls -lh "$MOUNT_POINT/" | sed 's/^/      /'
else
    mdir -i "$CIDATA_IMG" ::/ | sed 's/^/      /'
fi

# --- Finish ---
echo ""
if [ "$IMAGE_MODE" -eq 0 ]; then
    echo "[6/6] Syncing and unmounting..."
    sync
    umount "$MOUNT_POINT"
    rmdir "$MOUNT_POINT"
    MOUNT_POINT=""
else
    echo "[6/6] Writing CIDATA filesystem into $IMAGE_FILE at sector $PART_START..."
    dd if="$CIDATA_IMG" of="$IMAGE_FILE" bs=1M conv=notrunc,sparse \
        oflag=seek_bytes seek=$((PART_START * SECTOR_SIZE)) status=none
    rm -f "$CIDATA_IMG"
    echo "    Partition table:"
    sfdisk -l "$IMAGE_FILE" 2>/dev/null | sed 's/^/      /'
fi

echo ""
echo "============================================"
echo "USB build complete!"
echo "============================================"
echo ""
if [ "$IMAGE_MODE" -eq 1 ]; then
    echo "Image: $IMAGE_FILE"
    echo "  Write to a stick with: sudo dd if=$IMAGE_FILE of=/dev/sdX bs=4M conv=fsync status=progress"
    echo "  or attach it to a VM as a USB disk (see tests/vm/README.md)."
else
    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"
fi
echo ""