#!/bin/bash
#
# test-vm.sh — Create a test VM to validate the PXE server setup
#
# This script:
#   1. Builds a CIDATA ISO with autoinstall config, packages, playbook, and webapp
#   2. Launches an Ubuntu 24.04 Server VM on the default libvirt network
#   3. The VM auto-installs, then runs the Ansible playbook on first boot
#
# Usage:
#   ./test-vm.sh /path/to/ubuntu-24.04-live-server-amd64.iso
#
# After install completes (~10-15 min), access via:
#   virsh console pxe-test       (serial console, always works)
#   ssh pxe@<dhcp-ip>            (check: virsh domifaddr pxe-test)
#
# To clean up:
#   ./test-vm.sh --destroy

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
VM_NAME="pxe-test"
VM_DISK="/var/lib/libvirt/images/${VM_NAME}.qcow2"
CIDATA_ISO="${SCRIPT_DIR}/.${VM_NAME}-cidata.iso"
VM_RAM=4096
VM_CPUS=2
VM_DISK_SIZE=40  # GB

# --- Handle --destroy flag ---
if [ "${1:-}" = "--destroy" ]; then
    echo "Destroying test environment..."
    virsh destroy "$VM_NAME" 2>/dev/null || true
    virsh undefine "$VM_NAME" 2>/dev/null || true
    virsh vol-delete "${VM_NAME}.qcow2" --pool default 2>/dev/null || true
    rm -f "$CIDATA_ISO"
    rm -f "/tmp/${VM_NAME}-vmlinuz" "/tmp/${VM_NAME}-initrd"
    echo "Done."
    exit 0
fi

# --- Validate Ubuntu ISO argument ---
UBUNTU_ISO="${1:-}"
if [ -z "$UBUNTU_ISO" ] || [ ! -f "$UBUNTU_ISO" ]; then
    echo "Usage: $0 /path/to/ubuntu-24.04-live-server-amd64.iso"
    echo ""
    echo "Download from: https://ubuntu.com/download/server"
    echo ""
    echo "Other commands:"
    echo "  $0 --destroy    Remove the test VM and network"
    exit 1
fi

echo "============================================"
echo "PXE Server Test VM Setup"
echo "============================================"
echo ""

# --- Step 1: Build CIDATA ISO ---
echo "[1/4] Building CIDATA ISO..."
CIDATA_DIR=$(mktemp -d)

# Autoinstall config
cp "$SCRIPT_DIR/autoinstall/user-data" "$CIDATA_DIR/user-data"
touch "$CIDATA_DIR/meta-data"

# Offline .deb packages
if [ -d "$SCRIPT_DIR/offline-packages" ]; then
    mkdir -p "$CIDATA_DIR/packages"
    cp "$SCRIPT_DIR/offline-packages/"*.deb "$CIDATA_DIR/packages/" 2>/dev/null || true
    echo "    Copied $(ls -1 "$CIDATA_DIR/packages/"*.deb 2>/dev/null | wc -l) .deb packages"
else
    echo "    WARNING: No offline-packages/ directory. Run download-packages.sh first."
fi

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

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

# Pip wheels
if [ -d "$SCRIPT_DIR/pip-wheels" ]; then
    cp -r "$SCRIPT_DIR/pip-wheels" "$CIDATA_DIR/pip-wheels"
    echo "    Copied pip-wheels/"
elif [ -d "$SCRIPT_DIR/offline-packages/pip-wheels" ]; then
    cp -r "$SCRIPT_DIR/offline-packages/pip-wheels" "$CIDATA_DIR/pip-wheels"
    echo "    Copied pip-wheels/ (from offline-packages/)"
fi

# WinPE boot files (wimboot, boot.wim, BCD, ipxe.efi, etc.)
if [ -d "$SCRIPT_DIR/boot-files" ]; then
    for bf in "$SCRIPT_DIR/boot-files"/*; do
        [ -f "$bf" ] && cp "$bf" "$CIDATA_DIR/"
    done
    echo "    Copied boot-files/ (wimboot, boot.wim, ipxe.efi, etc.)"
fi

# Boot tools
if [ -d "$SCRIPT_DIR/boot-tools" ]; then
    cp -r "$SCRIPT_DIR/boot-tools" "$CIDATA_DIR/boot-tools"
    echo "    Copied boot-tools/"
fi

# Generate the CIDATA ISO
genisoimage -output "$CIDATA_ISO" -volid CIDATA -joliet -rock "$CIDATA_DIR" 2>/dev/null
CIDATA_SIZE=$(du -sh "$CIDATA_ISO" | cut -f1)
echo "    CIDATA ISO: $CIDATA_ISO ($CIDATA_SIZE)"
rm -rf "$CIDATA_DIR"

# --- Step 2: Create VM disk ---
echo ""
echo "[2/4] Creating VM disk (${VM_DISK_SIZE}GB)..."
if virsh vol-info "$VM_NAME.qcow2" --pool default &>/dev/null; then
    echo "    Disk already exists. Destroy first with: $0 --destroy"
    exit 1
fi
virsh vol-create-as default "${VM_NAME}.qcow2" "${VM_DISK_SIZE}G" --format qcow2

# --- Step 3: Extract kernel/initrd from ISO ---
echo ""
echo "[3/4] Extracting kernel and initrd from ISO..."
KERNEL="/tmp/${VM_NAME}-vmlinuz"
INITRD="/tmp/${VM_NAME}-initrd"
7z e -o/tmp -y "$UBUNTU_ISO" casper/vmlinuz casper/initrd 2>/dev/null
mv /tmp/vmlinuz "$KERNEL"
mv /tmp/initrd "$INITRD"
echo "    Extracted vmlinuz and initrd from casper/"

# --- Step 4: Launch VM ---
echo ""
echo "[4/4] Launching VM ($VM_NAME)..."

# Use the default libvirt network (NAT, 192.168.122.0/24) for install access.
# If br-pxe bridge exists, add a second NIC for the isolated PXE switch.
# The Ansible playbook will configure 10.9.100.1/24 on the PXE interface.
PXE_BRIDGE_ARGS=""
if ip link show br-pxe &>/dev/null; then
    PXE_BRIDGE_ARGS="--network bridge=br-pxe,model=virtio"
    echo "    Found br-pxe bridge, adding isolated switch NIC"
fi

virt-install \
    --name "$VM_NAME" \
    --memory "$VM_RAM" \
    --vcpus "$VM_CPUS" \
    --disk path="$VM_DISK",format=qcow2 \
    --disk path="$UBUNTU_ISO",device=cdrom,readonly=on \
    --disk path="$CIDATA_ISO",device=cdrom \
    --network network=default \
    $PXE_BRIDGE_ARGS \
    --os-variant ubuntu24.04 \
    --graphics none \
    --console pty,target_type=serial \
    --install kernel="$KERNEL",initrd="$INITRD",kernel_args="console=ttyS0,115200n8 autoinstall" \
    --noautoconsole

echo ""
echo "============================================"
echo "VM launched! The autoinstall will take ~10-15 minutes."
echo "============================================"
echo ""
echo "Watch progress:"
echo "  sudo virsh console $VM_NAME"
echo "  (Press Ctrl+] to detach)"
echo ""
echo "After install + first boot:"
echo "  Console: sudo virsh console $VM_NAME"
echo "  Find IP: sudo virsh domifaddr $VM_NAME"
echo "  SSH:     ssh pxe@<ip-from-above>"
echo ""
echo "NOTE: The Ansible playbook will change the VM's IP to 10.9.100.1."
echo "      After that, use 'virsh console' to access the VM."
echo "      On the VM, verify with: curl http://localhost:9009"
echo ""
echo "Manage:"
echo "  sudo virsh start $VM_NAME"
echo "  sudo virsh shutdown $VM_NAME"
echo "  $0 --destroy   (remove everything)"
echo ""