#!/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/, site-content/) 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/
#   --skip-site-content do not copy site-content/ (Office, Adobe, certs,
#                       Engineering Layer...: scripts/capture-site-content.sh).
#                       Without it a new server cannot build standard or
#                       engineer until that content is copied over by hand.
#   --site-vars FILE    put FILE on CIDATA as playbook/site-vars.yml
#                       (per-site values; see playbook/site-vars.example.yml)
#   --no-site-vars      build without site vars (defaults; the first boot
#                       picks the imaging NIC by its link). One of the two
#                       is required; playbook/site-vars.yml in the repo is
#                       never copied implicitly.
#   --yes               do not ask for confirmation (device mode still asks
#                       YES for a non-removable device)
#
# Refuses to build when a payload does not match its pin: the iPXE files in
# boot-files/ against the playbook's ipxe_*_sha256, every wheel in
# pip-wheels/ against its own package's hashes in webapp/requirements.txt,
# and the ISO against the volume id offline-packages/ was resolved for.
#
# 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 | --no-site-vars, --skip-enrollment,"
    echo "           --skip-bios, --skip-drivers, --skip-boot-tools, --skip-site-content, --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
SKIP_SITE_CONTENT=0
ASSUME_YES=0
SITE_VARS=""
NO_SITE_VARS=0
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 ;;
        --no-site-vars) NO_SITE_VARS=1; shift ;;
        --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 ;;
        --skip-site-content) SKIP_SITE_CONTENT=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"
SITE_CONTENT_DIR="$PAYLOAD_ROOT/site-content"
[ "$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=""
[ "$SKIP_SITE_CONTENT" -eq 1 ] && SITE_CONTENT_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

# The installer patches and installs from packages/ as an apt repo
# (autoinstall late-commands -> offline-apt-install.sh): it needs the index
# and the explicit list that download-packages.sh writes with the debs.
for f in "$OFFLINE_PKG_DIR/Packages" "$OFFLINE_PKG_DIR/explicit.list" "$REPO_ROOT/scripts/offline-apt-install.sh"; do
    if [ ! -f "$f" ]; then
        echo "ERROR: $f not found. Build the offline set with"
        echo "       scripts/download-packages.sh --iso $ISO_PATH $OFFLINE_PKG_DIR"
        exit 1
    fi
done

if [ -n "$SITE_VARS" ] && [ "$NO_SITE_VARS" -eq 1 ]; then
    echo "ERROR: give --site-vars FILE or --no-site-vars, not both"
    exit 1
fi
if [ -z "$SITE_VARS" ] && [ "$NO_SITE_VARS" -eq 0 ]; then
    echo "ERROR: give --site-vars FILE (per-site values, see playbook/site-vars.example.yml)"
    echo "       or --no-site-vars (the first boot then picks the imaging NIC itself and uses"
    echo "       the defaults for everything else)."
    exit 1
fi
if [ -n "$SITE_VARS" ]; then
    if [ ! -f "$SITE_VARS" ]; then
        echo "ERROR: --site-vars file not found: $SITE_VARS"
        exit 1
    fi
    # Both MACs are optional: an empty pxe_nic_mac makes the first boot pick
    # the imaging NIC by its link (only the imaging switch plugged in), an
    # empty corp_nic_mac makes every other port a DHCP client. One that is
    # set must be a MAC.
    for v in pxe_nic_mac corp_nic_mac; do
        line=$(grep -E "^$v:" "$SITE_VARS" | tail -1)
        [ -n "$line" ] || continue
        if grep -qE "^$v:[[:space:]]*([\"'][\"'])?[[:space:]]*(#.*)?$" <<<"$line"; then
            continue
        fi
        if ! grep -qE "^$v:[[:space:]]*[\"']?([0-9A-Fa-f]{2}:){5}[0-9A-Fa-f]{2}[\"']?[[:space:]]*(#.*)?$" <<<"$line"; then
            echo "ERROR: $SITE_VARS sets $v to something that is not a MAC (aa:bb:cc:dd:ee:ff) or empty"
            exit 1
        fi
    done
fi

# --- Payload pins: refuse to build a USB the server would reject ---
# iPXE chain: the same sha256 pins the playbook checks (single source: the
# ipxe_*_sha256 vars in playbook/pxe_server_setup.yml).
pin_of() {
    sed -n "s/^[[:space:]]*$1:[[:space:]]*\"\([0-9a-f]\{64\}\)\".*/\1/p" "$PLAYBOOK_DIR/pxe_server_setup.yml" | head -1
}
sha_ok() {   # sha_ok FILE PIN
    [ -f "$1" ] && [ "$(sha256sum <"$1" | cut -d' ' -f1)" = "$2" ]
}
for var in ipxe_tarball_sha256 ipxe_shim_sha256 ipxe_snponly_sha256 ipxe_bios_sha256; do
    [ -n "$(pin_of "$var")" ] || { echo "ERROR: no $var pin in playbook/pxe_server_setup.yml"; exit 1; }
done
if ! sha_ok "$BOOT_FILES_DIR/ipxeboot.tar.gz" "$(pin_of ipxe_tarball_sha256)"; then
    echo "ERROR: $BOOT_FILES_DIR/ipxeboot.tar.gz is missing or does not match ipxe_tarball_sha256"
    echo "       (scripts/fetch-ipxe.sh fetches the pinned release)"
    exit 1
fi
if ! sha_ok "$BOOT_FILES_DIR/ipxe-shimx64.efi" "$(pin_of ipxe_shim_sha256)"; then
    echo "ERROR: $BOOT_FILES_DIR/ipxe-shimx64.efi is missing or does not match ipxe_shim_sha256"
    exit 1
fi
for pair in "ipxeboot/x86_64-sb/snponly.efi ipxe_snponly_sha256" "ipxeboot/x86_64/ipxe.pxe ipxe_bios_sha256"; do
    member=${pair% *}; var=${pair#* }
    got=$(tar -xzOf "$BOOT_FILES_DIR/ipxeboot.tar.gz" "$member" 2>/dev/null | sha256sum | cut -d' ' -f1)
    if [ "$got" != "$(pin_of "$var")" ]; then
        echo "ERROR: $member in ipxeboot.tar.gz does not match $var"
        exit 1
    fi
done
echo "iPXE chain matches the playbook pins"

# Wheels: every wheel must be the pinned version of its package with one
# of THAT package's hashes, and each pinned package needs exactly one. No
# download fallback: the playbook refuses anything else.
if ! compgen -G "$PIP_WHEELS_DIR/*.whl" >/dev/null; then
    echo "ERROR: no wheels in $PIP_WHEELS_DIR; build them with scripts/download-packages.sh"
    exit 1
fi
python3 - "$PIP_WHEELS_DIR" "$WEBAPP_DIR/requirements.txt" <<'PY' || exit 1
import glob, hashlib, os, re, sys
wheel_dir, reqs = sys.argv[1], open(sys.argv[2]).read()
norm = lambda n: re.sub(r"[-_.]+", "_", n).lower()
pins, hashes = {}, {}
for m in re.finditer(r"^([A-Za-z0-9_.-]+)==([^\s\\]+)((?:\s*\\?\s*--hash=sha256:[0-9a-f]{64})+)", reqs, re.M):
    pins[norm(m.group(1))] = m.group(2)
    hashes[norm(m.group(1))] = set(re.findall(r"sha256:([0-9a-f]{64})", m.group(3)))
if not pins:
    sys.exit("ERROR: no hashed pins in webapp/requirements.txt")
seen, bad = {}, []
for whl in sorted(glob.glob(os.path.join(wheel_dir, "*.whl"))):
    base = os.path.basename(whl)
    name, ver = norm(base.split("-")[0]), base.split("-")[1]
    h = hashlib.sha256(open(whl, "rb").read()).hexdigest()
    if name not in pins:
        bad.append(f"{base}: {name} is not pinned")
    elif ver != pins[name]:
        bad.append(f"{base}: version {ver}, pinned {pins[name]}")
    elif h not in hashes[name]:
        bad.append(f"{base}: sha256 {h[:16]}... is not a pinned hash of {name}")
    seen.setdefault(name, []).append(base)
bad += [f"{n}: {len(seen.get(n, []))} wheels (need exactly 1)" for n in pins if len(seen.get(n, [])) != 1]
others = [os.path.basename(p) for p in glob.glob(os.path.join(wheel_dir, "*")) if not p.endswith(".whl")]
bad += [f"{o}: not a wheel" for o in others]
if bad:
    print("ERROR: pip-wheels/ does not match webapp/requirements.txt:")
    print("\n".join("       " + b for b in bad))
    sys.exit(1)
print(f"{len(seen)} wheels match their pinned hashes")
PY

# ISO: the offline set is resolved against one ISO's base system.
ISO_VOLID_NOW=$(xorriso -indev "$ISO_PATH" -pvd_info 2>/dev/null | sed -n 's/^Volume Id *: *//p' | head -1)
if [ -f "$OFFLINE_PKG_DIR/iso-volid" ]; then
    SET_VOLID=$(head -1 "$OFFLINE_PKG_DIR/iso-volid")
    if [ "$SET_VOLID" != "$ISO_VOLID_NOW" ]; then
        echo "ERROR: offline-packages/ was built for ISO '$SET_VOLID', this ISO is '$ISO_VOLID_NOW'."
        echo "       Rebuild it: scripts/download-packages.sh --iso $ISO_PATH $OFFLINE_PKG_DIR"
        exit 1
    fi
    echo "offline-packages/ matches ISO $ISO_VOLID_NOW"
else
    echo "WARNING: $OFFLINE_PKG_DIR/iso-volid missing (set built by an older download-packages.sh);"
    echo "         cannot check it was resolved against $ISO_VOLID_NOW."
fi

# site-content/ is restored file by file against its SHA256SUMS, so it is
# never split: a file FAT32 cannot hold refuses the build instead.
if [ -n "$SITE_CONTENT_DIR" ]; then
    if [ ! -d "$SITE_CONTENT_DIR" ]; then
        echo "ERROR: no site-content/ in $PAYLOAD_ROOT. Run scripts/capture-site-content.sh --apply"
        echo "       (or pass --skip-site-content: the new server then lacks Office, Adobe, certs"
        echo "       and the Engineering Layer until they are copied by hand)."
        exit 1
    fi
    if [ ! -f "$SITE_CONTENT_DIR/SHA256SUMS" ]; then
        echo "ERROR: site-content/SHA256SUMS missing. Re-run scripts/capture-site-content.sh --apply"
        exit 1
    fi
    SC_BIG=$(find -L "$SITE_CONTENT_DIR" -type f -size +4194303k)
    if [ -n "$SC_BIG" ]; then
        echo "ERROR: site-content/ has files of 4 GiB or more (FAT32 cannot hold them):"
        printf '%s\n' "$SC_BIG" | sed 's/^/    /'
        exit 1
    fi
    SC_LISTED=$(grep -c . "$SITE_CONTENT_DIR/SHA256SUMS")
    SC_FOUND=$(find -L "$SITE_CONTENT_DIR" -type f ! -name SHA256SUMS | wc -l)
    if [ "$SC_LISTED" -ne "$SC_FOUND" ]; then
        echo "ERROR: site-content/ holds $SC_FOUND files but SHA256SUMS lists $SC_LISTED."
        echo "       Re-run scripts/capture-site-content.sh --apply"
        exit 1
    fi
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" "$SITE_CONTENT_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 -rL "$@" "$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 the offline package set (an apt repo: debs + Packages + explicit.list)
# into packages/, and the script the installer runs on it to the root.
t_mkdir packages
t_cp packages "$OFFLINE_PKG_DIR/Packages" "$OFFLINE_PKG_DIR/explicit.list"
DEB_COUNT=0
if compgen -G "$OFFLINE_PKG_DIR/*.deb" >/dev/null; then
    t_cp packages "$OFFLINE_PKG_DIR"/*.deb
    DEB_COUNT=$(find -L "$OFFLINE_PKG_DIR" -maxdepth 1 -name '*.deb' | wc -l)
fi
t_cp "" "$REPO_ROOT/scripts/offline-apt-install.sh"
echo "    Copied $DEB_COUNT .deb packages + Packages + explicit.list to packages/"

# Copy playbook directory, without a repo-local site-vars.yml: the site
# vars on the USB are exactly --site-vars (or none with --no-site-vars).
mkdir -p "$WORK/stage"
cp -rL "$PLAYBOOK_DIR" "$WORK/stage/playbook"
rm -f "$WORK/stage/playbook/site-vars.yml"
if [ -n "$SITE_VARS" ]; then
    cp "$SITE_VARS" "$WORK/stage/playbook/site-vars.yml"
fi
t_cp "" "$WORK/stage/playbook"
rm -rf "$WORK/stage"
echo "    Copied playbook/"
if [ -n "$SITE_VARS" ]; then
    echo "    Copied $SITE_VARS as playbook/site-vars.yml"
else
    echo "    WARNING: --no-site-vars: the first-boot playbook stops until"
    echo "             pxe_nic_mac / corp_nic_mac are set on the server."
fi
# The unattend linter the playbook runs before it replaces the answer files
# ({{ usb_mount }}/../scripts/lint-unattend.py).
t_mkdir scripts
t_cp scripts "$REPO_ROOT/scripts/lint-unattend.py"

# 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 the pinned pip wheels (checked against requirements.txt above)
t_cp "" "$PIP_WHEELS_DIR"
WHEEL_COUNT=$(find -L "$PIP_WHEELS_DIR" -name '*.whl' | wc -l)
echo "    Copied pip-wheels/ ($WHEEL_COUNT wheels)"

# 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) 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 -L "$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 the site content (capture-site-content.sh); the playbook checks
# SHA256SUMS before it restores anything.
if [ -n "$SITE_CONTENT_DIR" ]; then
    t_cp "" "$SITE_CONTENT_DIR"
    echo "    Copied site-content/ ($(du -shL "$SITE_CONTENT_DIR" | cut -f1), $SC_FOUND files)"
else
    echo "    Skipped site-content/ (--skip-site-content)"
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 -L "$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..."
    # FAT32 cannot hold a file of 4 GiB or more, and the playbook's import
    # (cp -rn) does not reassemble split parts, so such a file stops here.
    BIG=$(find -L "$WINPE_SOURCE" -type f -size +4194303k 2>/dev/null)
    if [ -n "$BIG" ]; then
        echo "ERROR: WinPE content has files of 4 GiB or more, which FAT32 (CIDATA) cannot hold:"
        printf '%s\n' "$BIG" | sed 's/^/       /'
        echo "       Import those images through the webapp (Images > Import) instead."
        exit 1
    fi
    t_mkdir images
    t_cp images "$WINPE_SOURCE"/*
    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. Cable both NICs: the imaging NIC (pxe_nic_mac) to the isolated"
    echo "     imaging switch, the corp NIC (corp_nic_mac) to the GE network"
    echo "  2. Insert USB into the target machine, boot from it (F12 / boot menu)"
    echo "  3. Ubuntu auto-installs; on first boot the playbook pins the NICs by"
    echo "     MAC (imaging = static 172.16.9.1, corp = DHCP) and sets up PXE"
    echo "  4. Check /var/log/first-boot.log; a failure shows a banner at login"
    echo "     and on the console, and first boot retries at the next reboot"
fi
echo ""