#!/bin/bash
#
# test-lab.sh — Full PXE lab: server + client VMs on an isolated network
#
# Creates an isolated libvirt network, boots the PXE server from the
# Proxmox installer ISO, then launches a UEFI PXE client to test the
# full boot chain (DHCP -> TFTP -> iPXE -> boot menu).
#
# Usage:
#   ./test-lab.sh /path/to/ubuntu-24.04.iso         # Launch server
#   ./test-lab.sh --client                           # Launch PXE client
#   ./test-lab.sh --status                           # Check if server is ready
#   ./test-lab.sh --destroy                          # Remove everything
#
# Workflow:
#   1. Run the script with the Ubuntu ISO — server VM starts installing
#   2. Wait ~15 minutes (monitor with: sudo virsh console pxe-lab-server)
#   3. Run --status to check if PXE services are up
#   4. Run --client to launch a PXE client VM
#   5. Open virt-viewer to watch the client PXE boot:
#        virt-viewer pxe-lab-client

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
NET_NAME="pxe-lab"
SERVER_NAME="pxe-lab-server"
CLIENT_NAME="pxe-lab-client"
SERVER_DISK="/var/lib/libvirt/images/${SERVER_NAME}.qcow2"
PROXMOX_ISO="$SCRIPT_DIR/pxe-server-proxmox.iso"
VM_RAM=4096
VM_CPUS=2
VM_DISK_SIZE=40  # GB

# --- Helper: check if we can run virsh ---
check_virsh() {
    if ! virsh net-list &>/dev/null; then
        echo "ERROR: Cannot connect to libvirt. Are you in the 'libvirt' group?"
        echo "  sudo usermod -aG libvirt $USER && newgrp libvirt"
        exit 1
    fi
}

# --- Helper: ensure network exists ---
ensure_network() {
    if ! virsh net-info "$NET_NAME" &>/dev/null; then
        echo "Creating isolated network ($NET_NAME)..."
        NET_XML=$(mktemp)
        cat > "$NET_XML" << EOF
<network>
  <name>$NET_NAME</name>
  <bridge name="virbr-pxe" stp="on" delay="0"/>
</network>
EOF
        virsh net-define "$NET_XML" >/dev/null
        rm "$NET_XML"
    fi
    if ! virsh net-info "$NET_NAME" 2>/dev/null | grep -q "Active:.*yes"; then
        virsh net-start "$NET_NAME" >/dev/null
    fi
}

# =====================================================================
# --destroy: Remove everything
# =====================================================================
if [ "${1:-}" = "--destroy" ]; then
    check_virsh
    echo "Destroying PXE lab environment..."
    virsh destroy "$CLIENT_NAME" 2>/dev/null || true
    virsh undefine "$CLIENT_NAME" --nvram 2>/dev/null || true
    virsh vol-delete "${CLIENT_NAME}.qcow2" --pool default 2>/dev/null || true
    virsh destroy "$SERVER_NAME" 2>/dev/null || true
    virsh undefine "$SERVER_NAME" 2>/dev/null || true
    virsh vol-delete "${SERVER_NAME}.qcow2" --pool default 2>/dev/null || true
    rm -f "/tmp/${SERVER_NAME}-vmlinuz" "/tmp/${SERVER_NAME}-initrd"
    virsh net-destroy "$NET_NAME" 2>/dev/null || true
    virsh net-undefine "$NET_NAME" 2>/dev/null || true
    echo "Done."
    exit 0
fi

# =====================================================================
# --status: Check if PXE server is ready
# =====================================================================
if [ "${1:-}" = "--status" ]; then
    check_virsh
    echo "PXE Lab Status"
    echo "============================================"

    # Network
    if virsh net-info "$NET_NAME" &>/dev/null; then
        echo "  Network ($NET_NAME): $(virsh net-info "$NET_NAME" 2>/dev/null | grep Active | awk '{print $2}')"
    else
        echo "  Network ($NET_NAME): not defined"
    fi

    # Server VM
    if virsh dominfo "$SERVER_NAME" &>/dev/null; then
        STATE=$(virsh domstate "$SERVER_NAME" 2>/dev/null)
        echo "  Server ($SERVER_NAME): $STATE"
    else
        echo "  Server ($SERVER_NAME): not defined"
    fi

    # Client VM
    if virsh dominfo "$CLIENT_NAME" &>/dev/null; then
        STATE=$(virsh domstate "$CLIENT_NAME" 2>/dev/null)
        echo "  Client ($CLIENT_NAME): $STATE"
    else
        echo "  Client ($CLIENT_NAME): not defined"
    fi

    # Try to check PXE services on the server
    # Add a temporary IP to the bridge so we can reach the server
    echo ""
    echo "Checking PXE server services..."
    BRIDGE_HAS_IP=false
    ADDED_IP=false
    if ip addr show virbr-pxe 2>/dev/null | grep -q "10.9.100.254"; then
        BRIDGE_HAS_IP=true
    else
        # Need sudo for IP manipulation — try it, skip if unavailable
        if sudo -n ip addr add 10.9.100.254/24 dev virbr-pxe 2>/dev/null; then
            BRIDGE_HAS_IP=true
            ADDED_IP=true
        fi
    fi

    if [ "$BRIDGE_HAS_IP" = true ]; then
        # Check each service with a short timeout
        for check in \
            "DHCP/TFTP (dnsmasq):10.9.100.1:69:udp" \
            "HTTP (Apache):10.9.100.1:80:tcp" \
            "iPXE boot script:10.9.100.1:4433:tcp" \
            "Samba:10.9.100.1:445:tcp" \
            "Webapp:10.9.100.1:9009:tcp"; do
            LABEL="${check%%:*}"
            REST="${check#*:}"
            HOST="${REST%%:*}"
            REST="${REST#*:}"
            PORT="${REST%%:*}"
            PROTO="${REST#*:}"

            if [ "$PROTO" = "tcp" ]; then
                if timeout 2 bash -c "echo >/dev/tcp/$HOST/$PORT" 2>/dev/null; then
                    echo "  [UP]   $LABEL (port $PORT)"
                else
                    echo "  [DOWN] $LABEL (port $PORT)"
                fi
            else
                # UDP — just check if host is reachable
                if ping -c1 -W1 "$HOST" &>/dev/null; then
                    echo "  [PING] $LABEL (host reachable)"
                else
                    echo "  [DOWN] $LABEL (host unreachable)"
                fi
            fi
        done

        # Clean up the temporary IP (only if we added it)
        if [ "$ADDED_IP" = true ]; then
            sudo -n ip addr del 10.9.100.254/24 dev virbr-pxe 2>/dev/null || true
        fi
    else
        echo "  (Cannot reach server — bridge not available)"
        echo "  Use 'sudo virsh console $SERVER_NAME' to check manually"
    fi

    echo ""
    echo "Commands:"
    echo "  sudo virsh console $SERVER_NAME    # Server serial console"
    echo "  virt-viewer $CLIENT_NAME           # Client VNC display"
    exit 0
fi

# =====================================================================
# --client: Launch PXE client VM
# =====================================================================
if [ "${1:-}" = "--client" ]; then
    check_virsh
    ensure_network

    # Check if server VM exists and is running
    if ! virsh domstate "$SERVER_NAME" 2>/dev/null | grep -q "running"; then
        echo "WARNING: Server VM ($SERVER_NAME) is not running."
        echo "  The PXE client needs the server for DHCP and boot files."
        read -rp "  Continue anyway? (y/N): " PROCEED
        if [[ ! "$PROCEED" =~ ^[Yy]$ ]]; then
            exit 1
        fi
    fi

    # Remove existing client if present
    if virsh dominfo "$CLIENT_NAME" &>/dev/null; then
        echo "Removing existing client VM..."
        virsh destroy "$CLIENT_NAME" 2>/dev/null || true
        virsh undefine "$CLIENT_NAME" --nvram 2>/dev/null || true
        virsh vol-delete "${CLIENT_NAME}.qcow2" --pool default 2>/dev/null || true
    fi

    echo "Launching PXE client ($CLIENT_NAME)..."
    echo "  UEFI PXE boot on network: $NET_NAME"
    echo ""

    virsh vol-create-as default "${CLIENT_NAME}.qcow2" "${VM_DISK_SIZE}G" --format qcow2 >/dev/null

    virt-install \
        --name "$CLIENT_NAME" \
        --memory "$VM_RAM" \
        --vcpus "$VM_CPUS" \
        --disk "vol=default/${CLIENT_NAME}.qcow2" \
        --network network="$NET_NAME",model=virtio \
        --os-variant ubuntu24.04 \
        --boot uefi,network \
        --graphics vnc,listen=0.0.0.0 \
        --noautoconsole

    # Get VNC port
    VNC_PORT=$(virsh vncdisplay "$CLIENT_NAME" 2>/dev/null | sed 's/://' || echo "?")
    VNC_PORT=$((5900 + VNC_PORT))

    echo ""
    echo "============================================"
    echo "PXE client launched!"
    echo "============================================"
    echo ""
    echo "Watch the PXE boot:"
    echo "  virt-viewer $CLIENT_NAME"
    echo "  (or VNC to localhost:$VNC_PORT)"
    echo ""
    echo "Expected boot sequence:"
    echo "  1. UEFI firmware -> PXE boot"
    echo "  2. DHCP from server (10.9.100.x)"
    echo "  3. TFTP download ipxe.efi"
    echo "  4. iPXE loads boot menu from port 4433"
    echo "  5. GE Aerospace PXE Boot Menu appears"
    echo ""
    echo "Manage:"
    echo "  sudo virsh reboot $CLIENT_NAME    # Retry PXE boot"
    echo "  sudo virsh destroy $CLIENT_NAME   # Stop client"
    echo "  $0 --destroy                 # Remove everything"
    exit 0
fi

# =====================================================================
# Default: Launch PXE server VM
# =====================================================================
check_virsh

UBUNTU_ISO="${1:-}"
if [ -z "$UBUNTU_ISO" ] || [ ! -f "$UBUNTU_ISO" ]; then
    echo "Usage: sudo $0 /path/to/ubuntu-24.04-live-server-amd64.iso"
    echo ""
    echo "Commands:"
    echo "  $0 /path/to/ubuntu.iso    Launch PXE server VM"
    echo "  $0 --client               Launch PXE client VM"
    echo "  $0 --status               Check server readiness"
    echo "  $0 --destroy              Remove everything"
    exit 1
fi

# Check if Proxmox ISO exists, build if not
if [ ! -f "$PROXMOX_ISO" ]; then
    echo "Proxmox ISO not found. Building it first..."
    echo ""
    "$SCRIPT_DIR/build-proxmox-iso.sh" "$UBUNTU_ISO"
    echo ""
fi

# Check server doesn't already exist
if virsh dominfo "$SERVER_NAME" &>/dev/null; then
    echo "ERROR: Server VM already exists. Destroy first with: sudo $0 --destroy"
    exit 1
fi

echo "============================================"
echo "PXE Lab Environment Setup"
echo "============================================"
echo ""

# --- Step 1: Create isolated network ---
echo "[1/3] Setting up isolated network ($NET_NAME)..."
ensure_network
echo "    Bridge: virbr-pxe (isolated, no host DHCP)"

# --- Step 2: Extract kernel/initrd for direct boot ---
echo "[2/3] Extracting kernel and initrd from ISO..."
KERNEL="/tmp/${SERVER_NAME}-vmlinuz"
INITRD="/tmp/${SERVER_NAME}-initrd"
7z e -o/tmp -y "$PROXMOX_ISO" casper/vmlinuz casper/initrd >/dev/null 2>&1
mv /tmp/vmlinuz "$KERNEL"
mv /tmp/initrd "$INITRD"
echo "    Extracted vmlinuz and initrd"

# --- Step 3: Launch server VM ---
echo "[3/3] Launching PXE server ($SERVER_NAME)..."

virsh vol-create-as default "${SERVER_NAME}.qcow2" "${VM_DISK_SIZE}G" --format qcow2 >/dev/null

virt-install \
    --name "$SERVER_NAME" \
    --memory "$VM_RAM" \
    --vcpus "$VM_CPUS" \
    --disk "vol=default/${SERVER_NAME}.qcow2" \
    --disk path="$PROXMOX_ISO",device=cdrom,readonly=on \
    --network network="$NET_NAME" \
    --os-variant ubuntu24.04 \
    --graphics none \
    --console pty,target_type=serial \
    --install kernel="$KERNEL",initrd="$INITRD",kernel_args="console=ttyS0,115200n8 autoinstall ds=nocloud\;s=/cdrom/server/" \
    --noautoconsole

echo ""
echo "============================================"
echo "PXE server VM launched!"
echo "============================================"
echo ""
echo "The autoinstall + first-boot will take ~15 minutes."
echo ""
echo "Step 1 — Monitor the server install:"
echo "  virsh console $SERVER_NAME"
echo "  (Ctrl+] to detach)"
echo ""
echo "Step 2 — Check when services are ready:"
echo "  $0 --status"
echo ""
echo "Step 3 — Launch a PXE client to test booting:"
echo "  $0 --client"
echo ""
echo "Cleanup:"
echo "  $0 --destroy"
echo ""