#!/bin/bash
#
# build-proxmox-iso.sh — Build a self-contained PXE server installer ISO for Proxmox
#
# Repackages the Ubuntu 24.04 Server ISO with:
#   - Autoinstall configuration (zero-touch install)
#   - All offline .deb packages and Python wheels
#   - Ansible playbook, Flask webapp, and boot tools
#
# The resulting ISO can be uploaded to Proxmox, attached to a VM, and booted.
# Ubuntu auto-installs, then first-boot configures all PXE services automatically.
#
# Usage:
#   ./build-proxmox-iso.sh /path/to/ubuntu-24.04-live-server-amd64.iso [output.iso]
#
# Prerequisites (on build workstation):
#   sudo apt install xorriso p7zip-full
#
# Before building, run:
#   ./download-packages.sh        (downloads offline .debs + pip wheels)
#   ./prepare-boot-tools.sh ...   (extracts Clonezilla, Blancco, Memtest)

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"
WEBAPP_DIR="$SCRIPT_DIR/webapp"
PIP_WHEELS_DIR="$SCRIPT_DIR/pip-wheels"
BOOT_TOOLS_DIR="$SCRIPT_DIR/boot-tools"

# --- Validate arguments ---
if [ $# -lt 1 ]; then
    echo "Usage: $0 /path/to/ubuntu-24.04-live-server-amd64.iso [output.iso]"
    echo ""
    echo "  Creates a self-contained ISO for deploying the PXE server in Proxmox."
    echo "  The ISO auto-installs Ubuntu and configures all PXE services."
    echo ""
    echo "Prerequisites:"
    echo "  sudo apt install xorriso p7zip-full"
    exit 1
fi

UBUNTU_ISO="$(realpath "$1")"
OUTPUT_ISO="${2:-$SCRIPT_DIR/pxe-server-proxmox.iso}"

# --- Validate prerequisites ---
echo "============================================"
echo "PXE Server Proxmox ISO Builder"
echo "============================================"
echo ""

MISSING_CMDS=()
for cmd in xorriso 7z; do
    if ! command -v "$cmd" &>/dev/null; then
        MISSING_CMDS+=("$cmd")
    fi
done

if [ ${#MISSING_CMDS[@]} -gt 0 ]; then
    echo "ERROR: Missing required tools: ${MISSING_CMDS[*]}"
    echo "Install with: sudo apt install xorriso p7zip-full"
    exit 1
fi

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

# Quick sanity check: ensure it looks like an Ubuntu Server ISO
ISO_CONTENTS=$(7z l "$UBUNTU_ISO" 2>&1) || true
if ! echo "$ISO_CONTENTS" | grep -q "casper/vmlinuz"; then
    echo "ERROR: Does not appear to be an Ubuntu Server ISO (missing casper/vmlinuz)"
    exit 1
fi

if [ ! -f "$AUTOINSTALL_DIR/user-data" ]; then
    echo "ERROR: user-data not found at $AUTOINSTALL_DIR/user-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 "Ubuntu ISO : $UBUNTU_ISO"
echo "Output ISO : $OUTPUT_ISO"
echo "Source Dir : $SCRIPT_DIR"
echo ""

# --- Setup work directory with cleanup trap ---
WORK_DIR=$(mktemp -d)
cleanup() { rm -rf "$WORK_DIR"; }
trap cleanup EXIT

EXTRACT_DIR="$WORK_DIR/iso"
BOOT_IMG_DIR="$WORK_DIR/BOOT"

# --- Step 1: Extract Ubuntu ISO ---
echo "[1/6] Extracting Ubuntu ISO..."
7z x -o"$EXTRACT_DIR" "$UBUNTU_ISO" -y >/dev/null 2>&1

# 7z extracts [BOOT] directory containing EFI images needed for rebuild
# Move it out so it doesn't end up in the final ISO filesystem
if [ -d "$EXTRACT_DIR/[BOOT]" ]; then
    mv "$EXTRACT_DIR/[BOOT]" "$BOOT_IMG_DIR"
    echo "    Extracted boot images for BIOS + UEFI"
else
    echo "ERROR: [BOOT] directory not found in extracted ISO"
    echo "    The Ubuntu ISO may be corrupted or an unsupported version."
    exit 1
fi

# Ensure files are writable (ISO extraction may set read-only)
chmod -R u+w "$EXTRACT_DIR"

# --- Step 2: Generate autoinstall user-data ---
echo "[2/6] Generating autoinstall configuration..."
mkdir -p "$EXTRACT_DIR/server"
touch "$EXTRACT_DIR/server/meta-data"

# Reuse the common sections (identity, network, storage, SSH) from existing user-data
# and replace late-commands with ISO-specific versions
sed '/^  late-commands:/,$d' "$AUTOINSTALL_DIR/user-data" > "$EXTRACT_DIR/server/user-data"

# Append ISO-specific late-commands
cat >> "$EXTRACT_DIR/server/user-data" << 'LATE_COMMANDS'
  late-commands:
    # Copy project files from ISO (/cdrom/pxe-data/) to the installed system
    - mkdir -p /target/opt/pxe-setup
    - cp -r /cdrom/pxe-data/packages /target/opt/pxe-setup/ 2>/dev/null || true
    - cp -r /cdrom/pxe-data/playbook /target/opt/pxe-setup/ 2>/dev/null || true
    - cp -r /cdrom/pxe-data/webapp /target/opt/pxe-setup/ 2>/dev/null || true
    - cp -r /cdrom/pxe-data/pip-wheels /target/opt/pxe-setup/ 2>/dev/null || true
    - cp -r /cdrom/pxe-data/boot-tools /target/opt/pxe-setup/ 2>/dev/null || true
    # Copy boot files (wimboot, boot.wim, BCD, ipxe.efi, etc.) from pxe-data root
    - sh -c 'for f in /cdrom/pxe-data/*; do [ -f "$f" ] && cp "$f" /target/opt/pxe-setup/; done' || true

    # Install deb packages in target chroot
    - |
      curtin in-target --target=/target -- bash -c '
      if compgen -G "/opt/pxe-setup/packages/*.deb" > /dev/null; then
        dpkg -i /opt/pxe-setup/packages/*.deb 2>/dev/null || true
        dpkg -i /opt/pxe-setup/packages/*.deb 2>/dev/null || true
        if command -v nmcli >/dev/null; then
          systemctl enable NetworkManager
        fi
      fi
      '

    # Create first-boot script (reads from local /opt/pxe-setup/)
    - |
      curtin in-target --target=/target -- bash -c '
      cat <<"EOF" > /opt/first-boot.sh
      #!/bin/bash
      SRC=/opt/pxe-setup
      # Install all offline .deb packages
      if compgen -G "$SRC/packages/*.deb" > /dev/null; then
        dpkg -i $SRC/packages/*.deb 2>/dev/null || true
        dpkg -i $SRC/packages/*.deb 2>/dev/null || true
      fi
      # Run the Ansible playbook (override USB paths to local source)
      if [ -f $SRC/playbook/pxe_server_setup.yml ]; then
        cd $SRC/playbook
        ansible-playbook -i localhost, -c local pxe_server_setup.yml \
          -e usb_root=$SRC -e usb_mount=$SRC/playbook
      fi
      # Disable rc.local to prevent rerunning
      sed -i "s|^/opt/first-boot.sh.*|# &|" /etc/rc.local
      lvextend -r -l +100%FREE /dev/mapper/ubuntu--vg-ubuntu--lv || true
      # Clean up large setup files to save disk space
      rm -rf $SRC/packages $SRC/pip-wheels $SRC/boot-tools
      rm -f $SRC/boot.wim $SRC/boot.sdi $SRC/bootx64.efi $SRC/wimboot $SRC/ipxe.efi $SRC/BCD $SRC/boot.stl
      EOF
      '
    - curtin in-target --target=/target -- chmod +x /opt/first-boot.sh

    # Create rc.local to run first-boot on next startup
    - |
      curtin in-target --target=/target -- bash -c '
      cat <<"EOF" > /etc/rc.local
      #!/bin/bash
      /opt/first-boot.sh > /var/log/first-boot.log 2>&1 &
      exit 0
      EOF
      '
    - curtin in-target --target=/target -- chmod +x /etc/rc.local

  user-data:
    disable_root: false

  refresh-installer:
    update: no
LATE_COMMANDS

echo "    Generated server/user-data and server/meta-data"

# --- Step 3: Copy project files to pxe-data/ ---
echo "[3/6] Copying project files to ISO..."
PXE_DATA="$EXTRACT_DIR/pxe-data"
mkdir -p "$PXE_DATA"

# Offline .deb packages
if [ -d "$OFFLINE_PKG_DIR" ]; then
    mkdir -p "$PXE_DATA/packages"
    DEB_COUNT=0
    for deb in "$OFFLINE_PKG_DIR"/*.deb; do
        if [ -f "$deb" ]; then
            cp "$deb" "$PXE_DATA/packages/"
            DEB_COUNT=$((DEB_COUNT + 1))
        fi
    done
    echo "    Copied $DEB_COUNT .deb packages"
else
    echo "    WARNING: No offline-packages/ directory. Run download-packages.sh first."
fi

# Ansible playbook
mkdir -p "$PXE_DATA/playbook"
cp "$PLAYBOOK_DIR/"* "$PXE_DATA/playbook/" 2>/dev/null || true
echo "    Copied playbook/"

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

# Python 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?)"
        rmdir "$PIP_WHEELS_DIR" 2>/dev/null || true
    fi
fi
if [ -d "$PIP_WHEELS_DIR" ]; then
    cp -r "$PIP_WHEELS_DIR" "$PXE_DATA/pip-wheels"
    WHEEL_COUNT=$(find "$PIP_WHEELS_DIR" -name '*.whl' | wc -l)
    echo "    Copied pip-wheels/ ($WHEEL_COUNT wheels)"
fi

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

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

# --- Step 4: Modify GRUB for autoinstall ---
echo "[4/6] Configuring autoinstall boot..."
GRUB_CFG="$EXTRACT_DIR/boot/grub/grub.cfg"

if [ ! -f "$GRUB_CFG" ]; then
    echo "ERROR: boot/grub/grub.cfg not found in extracted ISO"
    exit 1
fi

# Add autoinstall kernel parameter with nocloud datasource pointing to /cdrom/server/
# The semicolon must be escaped as \; in GRUB (it's a command separator)
# Apply to both regular and HWE kernels
sed -i 's|/casper/vmlinuz\b|/casper/vmlinuz autoinstall ds=nocloud\\;s=/cdrom/server/|g' "$GRUB_CFG"
sed -i 's|/casper/hwe-vmlinuz\b|/casper/hwe-vmlinuz autoinstall ds=nocloud\\;s=/cdrom/server/|g' "$GRUB_CFG"

# Reduce timeout for automatic boot (1 second instead of default 30)
sed -i 's/set timeout=.*/set timeout=1/' "$GRUB_CFG"

echo "    Modified GRUB: autoinstall enabled, timeout=1s"

# --- Step 5: Rebuild ISO ---
echo "[5/6] Rebuilding ISO (this may take a few minutes)..."

# Verify required boot images exist
EFI_IMG="$BOOT_IMG_DIR/2-Boot-NoEmul.img"
if [ ! -f "$EFI_IMG" ]; then
    echo "ERROR: EFI boot image not found at $EFI_IMG"
    exit 1
fi

if [ ! -f "$EXTRACT_DIR/boot/grub/i386-pc/eltorito.img" ]; then
    echo "ERROR: BIOS boot image not found at boot/grub/i386-pc/eltorito.img"
    exit 1
fi

xorriso -as mkisofs -r \
    -V 'PXE-SERVER' \
    -o "$OUTPUT_ISO" \
    --grub2-mbr --interval:local_fs:0s-15s:zero_mbrpt,zero_gpt:"$UBUNTU_ISO" \
    --protective-msdos-label \
    -partition_cyl_align off \
    -partition_offset 16 \
    --mbr-force-bootable \
    -append_partition 2 28732ac11ff8d211ba4b00a0c93ec93b "$EFI_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 4 -boot-info-table --grub2-boot-info \
    -eltorito-alt-boot \
    -e '--interval:appended_partition_2:::' \
    -no-emul-boot \
    "$EXTRACT_DIR"

# --- Step 6: Done ---
echo "[6/6] Cleaning up..."

ISO_SIZE=$(du -sh "$OUTPUT_ISO" | cut -f1)

echo ""
echo "============================================"
echo "Proxmox ISO build complete!"
echo "============================================"
echo ""
echo "Output: $OUTPUT_ISO ($ISO_SIZE)"
echo ""
echo "Proxmox deployment:"
echo "  1. Upload ISO to Proxmox storage (Datacenter -> Storage -> ISO Images)"
echo "  2. Create a new VM:"
echo "       - BIOS: OVMF (UEFI) — or SeaBIOS (both work)"
echo "       - Memory: 4096 MB"
echo "       - CPU: 2+ cores"
echo "       - Disk: 40+ GB (VirtIO SCSI)"
echo "       - Network: Bridge connected to isolated PXE network"
echo "  3. Attach ISO as CD-ROM and start the VM"
echo "  4. Ubuntu auto-installs (~10-15 minutes, zero interaction)"
echo "  5. After reboot, first-boot configures all PXE services"
echo "  6. Access webapp at http://10.9.100.1:9009"
echo ""
echo "NOTE: The VM's network bridge must be connected to your isolated PXE"
echo "      network. The server will use static IP 10.9.100.1/24."
echo ""