#!/usr/bin/env bash
# =============================================================================
# wds-deploy.sh — Package M-Archy archiso netboot artifacts for WDS + PXELinux
#
# PURPOSE:
#   Windows Deployment Services (WDS) on a Windows Server can serve PXE boot
#   images to clients via TFTP. This script:
#     1. Builds the M-Archy ISO + netboot tarball (by calling build.sh), or
#        reuses an existing tarball if --no-rebuild is given.
#     2. Extracts the kernel and initramfs from the netboot tarball.
#     3. Builds a TFTP directory tree compatible with WDS + PXELinux (syslinux).
#     4. Builds an HTTP directory tree for the airootfs.sfs squashfs root
#        (Arch's initrd fetches this over HTTP at boot time — TFTP is too slow).
#     5. Generates a pxelinux.cfg/default boot menu with multiple options.
#     6. Optionally zips the TFTP tree for easy transfer to a Windows Server.
#
# WHY WDS + PXELinux:
#   WDS natively speaks the Microsoft NBP (network boot program) protocol.
#   By configuring WDS to serve pxelinux.0 as the boot file, it becomes a
#   standards-compliant TFTP server that can chainload syslinux PXELinux,
#   which then presents our custom menu and loads the Arch kernel + initramfs.
#
# USAGE:
#   bash wds-deploy.sh --http-srv URL [OPTIONS] [OUT_DIR]
#
#   --http-srv URL      HTTP base URL where arch netboot files will be served
#                       (e.g. http://192.168.1.10/m-archy)
#                       Arch's initrd fetches airootfs.sfs over HTTP at boot time.
#                       Required.
#   --tftp-prefix PATH  Subdirectory within the WDS TFTP root for kernel/initramfs
#                       (default: m-archy)
#   --preconf [FILE]    Passed through to build.sh — embeds an answerfile into the ISO
#   --no-rebuild        Skip the archiso build if a netboot tarball already exists
#   OUT_DIR             Output directory (default: ~/m-archy-out)
#
# OUTPUT LAYOUT (inside OUT_DIR/wds-deploy/):
#   TFTP/               Copy contents to WDS TFTP root: C:\RemoteInstall\Boot\x64\
#   HTTP/               Serve as HTTP root at the URL given to --http-srv
#   wds-tftp.zip        Zip of TFTP/ — drop onto the Windows Server directly
#
# WDS DEPLOYMENT STEPS:
#   1. Serve HTTP/ over IIS/Nginx at the --http-srv URL
#   2. Extract wds-tftp.zip into C:\RemoteInstall\Boot\x64\
#   3. In WDS console → server Properties → Boot tab:
#        Set "Default boot program" for x64 to:  Boot\x64\pxelinux.0
#   4. PXE-boot a client — the M-Archy menu appears
# =============================================================================

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

# syslinux provides pxelinux.0 and the companion .c32 module files.
# On Arch Linux these live in /usr/lib/syslinux/bios/ after installing the
# syslinux package. WDS serves these to PXE-booting clients via TFTP.
SYSLINUX_BIOS="/usr/lib/syslinux/bios"

# ── Argument parsing ───────────────────────────────────────────────────────────
HTTP_SRV=""           # Required: base HTTP URL where airootfs.sfs is served
TFTP_PREFIX="m-archy" # Subdirectory under TFTP root for kernel + initramfs
PRECONF_ARGS=()       # Forwarded to build.sh if --preconf was given
NO_REBUILD=0          # 1 = skip build and reuse existing netboot tarball
OUT_ARG=""            # Positional output directory override

while [[ $# -gt 0 ]]; do
    case "$1" in
        --http-srv)
            [[ $# -gt 1 ]] || { echo "ERROR: --http-srv requires a URL" >&2; exit 1; }
            HTTP_SRV="$2"; shift 2 ;;
        --http-srv=*)
            HTTP_SRV="${1#--http-srv=}"; shift ;;
        --tftp-prefix)
            [[ $# -gt 1 ]] || { echo "ERROR: --tftp-prefix requires a value" >&2; exit 1; }
            TFTP_PREFIX="$2"; shift 2 ;;
        --tftp-prefix=*)
            TFTP_PREFIX="${1#--tftp-prefix=}"; shift ;;
        --preconf)
            # Forward --preconf [optional-file] to build.sh verbatim.
            if [[ $# -gt 1 && "${2:0:1}" != "-" ]]; then
                PRECONF_ARGS=(--preconf "$2"); shift 2
            else
                PRECONF_ARGS=(--preconf); shift
            fi ;;
        --preconf=*)
            # Forward --preconf=file to build.sh verbatim.
            PRECONF_ARGS=("$1"); shift ;;
        --no-rebuild)
            NO_REBUILD=1; shift ;;
        -*)
            echo "Unknown flag: $1" >&2; exit 1 ;;
        *)
            OUT_ARG="$1"; shift ;;
    esac
done

OUT_DIR="${OUT_ARG:-${OUT_DIR:-$HOME/m-archy-out}}"
# WDS_DIR is a subdirectory of OUT_DIR where we build the TFTP + HTTP trees.
WDS_DIR="$OUT_DIR/wds-deploy"

# ── Validate required arguments ───────────────────────────────────────────────
# --http-srv is mandatory because the Arch initrd must know where to download
# airootfs.sfs from at boot time. Without it, the kernel boots but the root
# filesystem cannot be mounted and the system hangs.
if [[ -z "$HTTP_SRV" ]]; then
    echo "ERROR: --http-srv <URL> is required." >&2
    echo "       The Arch initrd fetches airootfs.sfs over HTTP at boot." >&2
    echo "       Example: --http-srv http://192.168.1.10/m-archy" >&2
    exit 1
fi

# Strip trailing slash so URL concatenations are consistent throughout.
HTTP_SRV="${HTTP_SRV%/}"

# ── Ensure syslinux is available ──────────────────────────────────────────────
# pxelinux.0 is the PXE network boot program (NBP) that WDS sends to clients.
# Without syslinux installed we cannot produce the TFTP tree.
if [[ ! -f "$SYSLINUX_BIOS/pxelinux.0" ]]; then
    echo "syslinux not found — installing..."
    sudo pacman -S --noconfirm syslinux
fi

# Second check after attempted install — exit loudly if still missing.
if [[ ! -f "$SYSLINUX_BIOS/pxelinux.0" ]]; then
    echo "ERROR: $SYSLINUX_BIOS/pxelinux.0 still not found after install." >&2
    exit 1
fi

# ── Build ISO + netboot tarball (unless --no-rebuild) ─────────────────────────
# mkarchiso's 'netboot' build mode produces a .tar.gz alongside the .iso.
# The tarball contains the kernel, initramfs, and the squashfs tree needed
# for PXE netbooting. We extract kernel+initramfs to TFTP; the squashfs tree
# goes to the HTTP root.
NETBOOT_TARBALL="$(ls "$OUT_DIR/"*-netboot-*.tar.gz 2>/dev/null | head -n1 || true)"

if [[ "$NO_REBUILD" -eq 1 && -n "$NETBOOT_TARBALL" ]]; then
    echo "Skipping build — using existing tarball: $(basename "$NETBOOT_TARBALL")"
else
    echo "Building archiso (this may take a while)..."
    # Pass PRECONF_ARGS only when it is non-empty; the ${array[@]+"${array[@]}"}
    # idiom avoids "unbound variable" errors when the array is empty (set -u).
    bash "$SCRIPT_DIR/build.sh" "${PRECONF_ARGS[@]+"${PRECONF_ARGS[@]}"}" "$OUT_DIR"
    NETBOOT_TARBALL="$(ls "$OUT_DIR/"*-netboot-*.tar.gz 2>/dev/null | head -n1 || true)"
fi

[[ -n "$NETBOOT_TARBALL" ]] \
    || { echo "ERROR: No netboot tarball found in $OUT_DIR — build may have failed." >&2; exit 1; }

echo "Using netboot tarball: $(basename "$NETBOOT_TARBALL")"

# ── Extract netboot tarball ───────────────────────────────────────────────────
# We extract into a temporary directory, then pick out the specific files we
# need for each destination (TFTP vs HTTP) rather than serving the tarball
# layout directly.
EXTRACT_DIR="$OUT_DIR/.netboot-extracted"
rm -rf "$EXTRACT_DIR"
mkdir -p "$EXTRACT_DIR"
echo "Extracting netboot tarball..."
tar -xzf "$NETBOOT_TARBALL" -C "$EXTRACT_DIR"

# Locate the three key items inside the extracted tree.
# Using find + head -n1 handles variations in tarball layout across archiso versions.
VMLINUZ="$(find "$EXTRACT_DIR" -name 'vmlinuz-linux' | head -n1 || true)"
INITRAMFS="$(find "$EXTRACT_DIR" -name 'initramfs-linux.img' | head -n1 || true)"
ARCH_DIR="$(find "$EXTRACT_DIR" -maxdepth 2 -type d -name 'arch' | head -n1 || true)"

[[ -f "$VMLINUZ" ]]   || { echo "ERROR: vmlinuz-linux not found in netboot tarball." >&2; exit 1; }
[[ -f "$INITRAMFS" ]] || { echo "ERROR: initramfs-linux.img not found in netboot tarball." >&2; exit 1; }
[[ -d "$ARCH_DIR" ]]  || { echo "ERROR: arch/ directory not found in netboot tarball." >&2; exit 1; }

# ── Build WDS deployment tree ─────────────────────────────────────────────────
#
# TFTP root layout (mirrors what WDS serves via TFTP):
#   pxelinux.0                                    PXELinux bootloader
#   ldlinux.c32                                   required by pxelinux.0
#   menu.c32 + libcom32.c32 + libutil.c32         text boot menu modules
#   pxelinux.cfg/default                          boot menu configuration
#   <tftp-prefix>/arch/boot/x86_64/vmlinuz-linux       Linux kernel
#   <tftp-prefix>/arch/boot/x86_64/initramfs-linux.img  initrd (with archiso hooks)
#
# HTTP root layout (served at $HTTP_SRV/ over IIS/Nginx/Apache):
#   arch/x86_64/airootfs.sfs       the squashfs root filesystem fetched at boot
#   arch/x86_64/airootfs.sfs.sha512  integrity checksum
#   arch/pkglist.x86_64.txt        (and any other netboot metadata files)
#
# WHY two servers (TFTP + HTTP)?
#   TFTP is required by the PXE protocol for the initial NBP handshake and
#   kernel/initramfs delivery. However, TFTP is extremely slow for large files.
#   The airootfs.sfs (~500 MB squashfs) is downloaded by the initrd over HTTP,
#   which is orders of magnitude faster. This two-tier approach is the standard
#   Arch Linux netboot pattern.

TFTP_ROOT="$WDS_DIR/TFTP"
HTTP_ROOT="$WDS_DIR/HTTP"

rm -rf "$WDS_DIR"
mkdir -p \
    "$TFTP_ROOT/$TFTP_PREFIX/arch/boot/x86_64" \
    "$TFTP_ROOT/pxelinux.cfg" \
    "$HTTP_ROOT"

# ── Copy syslinux modules ─────────────────────────────────────────────────────
# pxelinux.0   — the actual NBP that WDS sends to the client via DHCP option 67.
# ldlinux.c32  — core syslinux library, required by pxelinux.0.
# menu.c32     — text-based interactive boot menu module.
# libcom32.c32 — common syslinux library needed by menu.c32.
# libutil.c32  — utility library needed by menu.c32.
echo "Copying syslinux PXELinux modules..."
REQUIRED_MODS=(pxelinux.0 ldlinux.c32 menu.c32 libcom32.c32 libutil.c32)
MISSING_MODS=()

for mod in "${REQUIRED_MODS[@]}"; do
    if [[ -f "$SYSLINUX_BIOS/$mod" ]]; then
        cp "$SYSLINUX_BIOS/$mod" "$TFTP_ROOT/"
    else
        # Older syslinux versions may not include all modules — warn but continue.
        # pxelinux.0 and ldlinux.c32 are essential; menu modules are optional.
        MISSING_MODS+=("$mod")
    fi
done

if [[ ${#MISSING_MODS[@]} -gt 0 ]]; then
    echo "Warning: missing syslinux modules (non-fatal for some setups): ${MISSING_MODS[*]}"
fi

# ── Copy kernel and initramfs ─────────────────────────────────────────────────
# These go into the TFTP tree because PXELinux loads them from the TFTP server
# as specified in pxelinux.cfg/default (KERNEL + initrd= directives).
echo "Copying kernel and initramfs..."
cp "$VMLINUZ"   "$TFTP_ROOT/$TFTP_PREFIX/arch/boot/x86_64/vmlinuz-linux"
cp "$INITRAMFS" "$TFTP_ROOT/$TFTP_PREFIX/arch/boot/x86_64/initramfs-linux.img"

# ── Copy HTTP content (airootfs squashfs and supporting files) ────────────────
# The entire arch/ directory from the netboot tarball goes to the HTTP root.
# The Arch initrd uses the archiso_http_srv and archisobasedir kernel parameters
# to locate and download airootfs.sfs from this path.
echo "Copying HTTP content (airootfs + supporting files)..."
cp -r "$ARCH_DIR" "$HTTP_ROOT/arch"

# ── Generate pxelinux.cfg/default ────────────────────────────────────────────
# This file is what PXELinux reads to build the boot menu. It must be at
# pxelinux.cfg/default relative to the TFTP root.
echo "Writing pxelinux.cfg/default..."

# Kernel path is relative to the TFTP root (PXELinux convention).
KERNEL_PATH="${TFTP_PREFIX}/arch/boot/x86_64/vmlinuz-linux"
INITRD_PATH="${TFTP_PREFIX}/arch/boot/x86_64/initramfs-linux.img"

# archiso kernel parameters:
#   archiso_http_srv=  — base URL of the HTTP server (must end with /)
#   archisobasedir=    — subdirectory within the HTTP server containing arch/x86_64/
#   ip=dhcp            — tell the initrd to acquire an IP via DHCP (required for HTTP)
APPEND_BASE="initrd=${INITRD_PATH} archiso_http_srv=${HTTP_SRV}/ archisobasedir=arch ip=dhcp"

cat > "$TFTP_ROOT/pxelinux.cfg/default" <<PXECFG
DEFAULT menu.c32
PROMPT 0
TIMEOUT 150
ONTIMEOUT m-archy

MENU TITLE M-Archy Arch Linux Installer

LABEL m-archy
  MENU LABEL M-Archy Arch Linux Installer
  KERNEL ${KERNEL_PATH}
  APPEND ${APPEND_BASE}

LABEL m-archy-console
  MENU LABEL M-Archy Arch Linux Installer (console only, no auto-launch)
  KERNEL ${KERNEL_PATH}
  APPEND ${APPEND_BASE} systemd.unit=multi-user.target

LABEL m-archy-ssh
  MENU LABEL M-Archy Arch Linux Installer (SSH access enabled)
  KERNEL ${KERNEL_PATH}
  APPEND ${APPEND_BASE} systemd.unit=multi-user.target sshd=1

LABEL local
  MENU LABEL Boot from local disk
  LOCALBOOT 0
PXECFG

# ── Zip TFTP directory for easy transfer to Windows Server ────────────────────
# Windows admins often prefer to receive a single zip rather than a directory
# tree via SCP. The zip can be extracted directly in C:\RemoteInstall\Boot\x64\.
ZIP_FILE="$WDS_DIR/wds-tftp.zip"
echo "Creating wds-tftp.zip..."
if command -v zip &>/dev/null; then
    # cd into TFTP_ROOT so the zip paths are relative (no leading path prefix).
    (cd "$TFTP_ROOT" && zip -r "$ZIP_FILE" .)
    echo "Created: $ZIP_FILE"
else
    echo "Note: 'zip' not installed — skipping wds-tftp.zip creation."
    echo "      Install with: sudo pacman -S zip"
fi

# ── Summary ───────────────────────────────────────────────────────────────────
echo
echo "========================================================================="
echo " WDS deployment package: $WDS_DIR"
echo "========================================================================="
echo
echo " TFTP/    → WDS TFTP root (copy to Windows Server)"
echo "   Destination: C:\\RemoteInstall\\Boot\\x64\\"
echo "   Tip: use wds-tftp.zip if created above, or robocopy/SCP the TFTP/ tree."
echo
echo " HTTP/    → HTTP root (serve at: ${HTTP_SRV}/)"
echo "   The initrd fetches: ${HTTP_SRV}/arch/x86_64/airootfs.sfs"
echo "   Serve with IIS, Nginx, or Apache pointing to the HTTP/ directory."
echo
echo " pxelinux.cfg/default kernel args summary:"
echo "   archiso_http_srv = ${HTTP_SRV}/"
echo "   archisobasedir   = arch"
echo
echo " WDS configuration steps:"
echo "   1. Copy TFTP/ contents to C:\\RemoteInstall\\Boot\\x64\\"
echo "   2. Serve HTTP/ over IIS/Nginx at: ${HTTP_SRV}/"
echo "   3. Open WDS console → right-click server → Properties → Boot tab:"
echo "        x64 default boot program: Boot\\x64\\pxelinux.0"
echo "        Check 'Always continue the PXE boot' for unknown clients"
echo "   4. If WDS manages DHCP, ensure option 67 is: Boot\\x64\\pxelinux.0"
echo "      If using a separate DHCP server, set:"
echo "        option 66 (next-server) = <WDS server IP>"
echo "        option 67 (boot-file)   = Boot\\x64\\pxelinux.0"
echo "   5. PXE-boot a client — the M-Archy menu should appear."
echo
echo " File sizes:"
du -sh "$TFTP_ROOT" "$HTTP_ROOT" 2>/dev/null | sed 's/^/   /'
echo "========================================================================="