#!/usr/bin/env bash
# wds-deploy.sh — Package M-Archy archiso netboot artifacts for WDS + PXELinux deployment
#
# 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_BIOS="/usr/lib/syslinux/bios"

# ── Argument parsing ───────────────────────────────────────────────────────────
HTTP_SRV=""
TFTP_PREFIX="m-archy"
PRECONF_ARGS=()
NO_REBUILD=0
OUT_ARG=""

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)
            if [[ $# -gt 1 && "${2:0:1}" != "-" ]]; then
                PRECONF_ARGS=(--preconf "$2"); shift 2
            else
                PRECONF_ARGS=(--preconf); shift
            fi ;;
        --preconf=*)
            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="$OUT_DIR/wds-deploy"

# ── Validate ──────────────────────────────────────────────────────────────────
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

HTTP_SRV="${HTTP_SRV%/}"  # strip trailing slash

# ── Ensure syslinux is available ──────────────────────────────────────────────
if [[ ! -f "$SYSLINUX_BIOS/pxelinux.0" ]]; then
    echo "syslinux not found — installing..."
    sudo pacman -S --noconfirm syslinux
fi

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) ─────────────────────────
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)..."
    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 ───────────────────────────────────────────────────
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"

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
#   pxelinux.cfg/default                          boot menu config
#   <tftp-prefix>/arch/boot/x86_64/vmlinuz-linux
#   <tftp-prefix>/arch/boot/x86_64/initramfs-linux.img
#
# HTTP root layout (served at $HTTP_SRV/ over IIS/Nginx/Apache):
#   arch/x86_64/airootfs.sfs                      fetched by initrd at boot
#   arch/x86_64/airootfs.sfs.sha512
#   arch/pkglist.x86_64.txt                       (and any other netboot files)

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 ─────────────────────────────────────────────────────
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
        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 ─────────────────────────────────────────────────
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) ────────────────
echo "Copying HTTP content (airootfs + supporting files)..."
cp -r "$ARCH_DIR" "$HTTP_ROOT/arch"

# ── Generate pxelinux.cfg/default ────────────────────────────────────────────
echo "Writing pxelinux.cfg/default..."

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

# archiso_http_srv must end with a slash; archisobasedir is the subdir within it
# that contains the x86_64/ squashfs tree
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 ────────────────────
ZIP_FILE="$WDS_DIR/wds-tftp.zip"
echo "Creating wds-tftp.zip..."
if command -v zip &>/dev/null; then
    (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 "========================================================================="