#!/bin/bash
# Tech to School — PXE VM Bootstrap Script
# Run this on a fresh Ubuntu 24.04 Server VM to set up the PXE boot stack
#
# Usage: ssh pxe 'bash -s' < scripts/pxe/pxe-vm-setup.sh
#    or: scp scripts/pxe/pxe-vm-setup.sh pxe:~ && ssh pxe 'chmod +x pxe-vm-setup.sh && sudo ./pxe-vm-setup.sh'

set -euo pipefail

PXE_DIR="/opt/pxe"
IPXE_VERSION="master"

echo "============================================"
echo "  PXE VM Bootstrap — Tech to School"
echo "============================================"
echo ""

# --- Check we're running as root ---
if [ "$EUID" -ne 0 ]; then
    echo "ERROR: Run as root (sudo ./pxe-vm-setup.sh)"
    exit 1
fi

# --- Step 1: System packages ---
echo "[1/6] Installing system packages..."
apt-get update -qq
apt-get install -y -qq \
    docker.io docker-compose-v2 \
    git build-essential liblzma-dev isolinux \
    curl wget net-tools tcpdump \
    wimtools

systemctl enable --now docker

# --- Step 2: Create directory structure ---
echo "[2/6] Creating PXE directory structure..."
mkdir -p $PXE_DIR/{tftpboot,winpe,images,scripts,deploy/{images,drivers,scripts}}

# --- Step 3: Build iPXE from source (UEFI + BIOS) ---
echo "[3/6] Building iPXE binaries..."
cd /tmp
if [ ! -d ipxe ]; then
    git clone --depth 1 https://github.com/ipxe/ipxe.git
fi
cd ipxe/src

# Embed script that chainloads our boot.ipxe from HTTP
cat > /tmp/embed.ipxe << 'IPXE'
#!ipxe
dhcp
chain http://${next-server}:8080/boot.ipxe || shell
IPXE

# Build UEFI binary (most modern devices)
echo "  Building ipxe.efi (UEFI)..."
make -j$(nproc) bin-x86_64-efi/ipxe.efi EMBED=/tmp/embed.ipxe 2>/dev/null
cp bin-x86_64-efi/ipxe.efi $PXE_DIR/tftpboot/

# Build BIOS binary (legacy fallback)
echo "  Building undionly.kpxe (BIOS)..."
make -j$(nproc) bin/undionly.kpxe EMBED=/tmp/embed.ipxe 2>/dev/null
cp bin/undionly.kpxe $PXE_DIR/tftpboot/

# --- Step 4: Download wimboot (iPXE Windows boot loader) ---
echo "[4/6] Downloading wimboot..."
WIMBOOT_URL="https://github.com/ipxe/wimboot/releases/latest/download/wimboot"
wget -q -O $PXE_DIR/tftpboot/wimboot "$WIMBOOT_URL"

# --- Step 5: Copy config files ---
echo "[5/6] Deploying configuration..."
# These should be scp'd from the MKL repo before running this script:
#   scp scripts/pxe/docker-compose.yml pxe:/opt/pxe/
#   scp scripts/pxe/nginx.conf pxe:/opt/pxe/
#   scp scripts/pxe/tftpboot/boot.ipxe pxe:/opt/pxe/tftpboot/
#   scp -r scripts/pxe/scripts/ pxe:/opt/pxe/

# Symlink deploy dirs so SMB share structure matches HTTP
ln -sf $PXE_DIR/images $PXE_DIR/deploy/images
ln -sf $PXE_DIR/scripts $PXE_DIR/deploy/scripts

# --- Step 6: Start Docker stack ---
echo "[6/6] Starting PXE services..."
cd $PXE_DIR
docker compose up -d

echo ""
echo "============================================"
echo "  PXE VM Bootstrap Complete!"
echo "============================================"
echo ""
echo "Directory structure:"
echo "  $PXE_DIR/"
echo "  ├── tftpboot/         # iPXE binaries (ipxe.efi, undionly.kpxe, wimboot)"
echo "  ├── winpe/            # WinPE boot files (BCD, boot.sdi, boot.wim)"
echo "  ├── images/           # WIM images (win11-hp-22h2.wim, etc.)"
echo "  ├── scripts/          # Deployment scripts (deploy.ps1, startnet.cmd)"
echo "  ├── deploy/           # SMB share root (symlinks to images/ and scripts/)"
echo "  ├── docker-compose.yml"
echo "  └── nginx.conf"
echo ""
echo "Next steps:"
echo "  1. Copy WinPE files to $PXE_DIR/winpe/ (BCD, boot.sdi, boot.wim)"
echo "     - Generate from Windows ADK: copype amd64 C:\WinPE"
echo "     - Customize startnet.cmd in the WIM, then copy BCD + boot.sdi + boot.wim"
echo "  2. Build and copy WIM images to $PXE_DIR/images/"
echo "  3. Add PXE options to Beast's dnsmasq config (next-server = this VM's IP)"
echo "  4. Test PXE boot on an HP 15-dy3001ds"
echo ""
echo "Services status:"
docker compose ps