#!/usr/bin/env bash
set -euo pipefail

# --- Configuration ---
LAB_DIR="/var/lib/harmony_pxe_test"
IMG_DIR="${LAB_DIR}/images"
STATE_DIR="${LAB_DIR}/state"
VM_OPN="opnsense-pxe"
VM_PXE="pxe-node-1"
NET_HARMONYLAN="harmonylan"

# Network settings for the isolated LAN
VLAN_CIDR="192.168.150.0/24"
VLAN_GW="192.168.150.1"
VLAN_MASK="255.255.255.0"

# VM Specifications
RAM_OPN="2048"
VCPUS_OPN="2"
DISK_OPN_GB="10"
OS_VARIANT_OPN="freebsd13.1" # Using a slightly more recent variant

RAM_PXE="4096"
VCPUS_PXE="2"
DISK_PXE_GB="40"
OS_VARIANT_LINUX="centos-stream9"

# ISO URLs and Paths
OPN_ISO="${IMG_DIR}/OPNsense-latest.iso"
CENTOS_ISO="${IMG_DIR}/CentOS-Stream-9-latest-boot.iso"
OPN_URL="https://mirror.wdc1.us.leaseweb.net/opnsense/releases/25.7/OPNsense-25.7-dvd-amd64.iso.bz2"
CENTOS_URL="https://mirror.stream.centos.org/9-stream/BaseOS/x86_64/os/images/boot.iso"

# Libvirt connection URI for system-wide daemon
CONNECT_URI="qemu:///system"

# --- Helper Functions ---
download_if_missing() {
  local url="$1"
  local dest="$2"
  if [[ ! -f "$dest" ]]; then
    echo "Downloading $url to $dest"
    mkdir -p "$(dirname "$dest")"
    local tmp
    tmp="$(mktemp)"
    curl -L --progress-bar "$url" -o "$tmp"
    case "$url" in
      *.bz2) bunzip2 -c "$tmp" > "$dest" && rm -f "$tmp" ;;
      *) mv "$tmp" "$dest" ;;
    esac
  else
    echo "Already present: $dest"
  fi
}

# Ensures a libvirt network is defined and active
ensure_network() {
  local net_name="$1"
  local net_xml_path="$2"
  if virsh --connect "${CONNECT_URI}" net-info "${net_name}" >/dev/null 2>&1; then
    echo "Network ${net_name} already exists."
  else
    echo "Defining network ${net_name} from ${net_xml_path}"
    virsh --connect "${CONNECT_URI}" net-define "${net_xml_path}"
  fi

  if ! virsh --connect "${CONNECT_URI}" net-info "${net_name}" | grep "Active: *yes"; then
    echo "Starting network ${net_name}..."
    virsh --connect "${CONNECT_URI}" net-start "${net_name}"
    virsh --connect "${CONNECT_URI}" net-autostart "${net_name}"
  fi
}

# Destroys a VM completely
destroy_vm() {
  local vm_name="$1"
  if virsh --connect "${CONNECT_URI}" dominfo "$vm_name" >/dev/null 2>&1; then
    echo "Destroying and undefining VM: ${vm_name}"
    virsh --connect "${CONNECT_URI}" destroy "$vm_name" || true
    virsh --connect "${CONNECT_URI}" undefine "$vm_name" --nvram
  fi
}

# Destroys a libvirt network
destroy_network() {
  local net_name="$1"
  if virsh --connect "${CONNECT_URI}" net-info "$net_name" >/dev/null 2>&1; then
    echo "Destroying and undefining network: ${net_name}"
    virsh --connect "${CONNECT_URI}" net-destroy "$net_name" || true
    virsh --connect "${CONNECT_URI}" net-undefine "$net_name"
  fi
}

# --- Main Logic ---
create_lab_environment() {
  # Create network definition files
  cat > "${STATE_DIR}/default.xml" <<EOF
<network>
  <name>default</name>
  <forward mode='nat'/>
  <bridge name='virbr0' stp='on' delay='0'/>
  <ip address='192.168.122.1' netmask='255.255.255.0'>
    <dhcp>
      <range start='192.168.122.100' end='192.168.122.200'/>
    </dhcp>
  </ip>
</network>
EOF

  cat > "${STATE_DIR}/${NET_HARMONYLAN}.xml" <<EOF
<network>
  <name>${NET_HARMONYLAN}</name>
  <forward mode='nat'/>
  <bridge name='virbr1' stp='on' delay='0'/>
  <ip address='${VLAN_GW}' netmask='${VLAN_MASK}'>
  </ip>
</network>
EOF

  # Ensure both networks exist and are active
  ensure_network "default" "${STATE_DIR}/default.xml"
  ensure_network "${NET_HARMONYLAN}" "${STATE_DIR}/${NET_HARMONYLAN}.xml"

  # --- Create OPNsense VM ---
  local disk_opn="${IMG_DIR}/${VM_OPN}.qcow2"
  if [[ ! -f "$disk_opn" ]]; then
    qemu-img create -f qcow2 "$disk_opn" "${DISK_OPN_GB}G"
  fi

  echo "Creating OPNsense VM..."
  virt-install \
    --connect "${CONNECT_URI}" \
    --name "${VM_OPN}" \
    --ram "${RAM_OPN}" \
    --vcpus "${VCPUS_OPN}" \
    --cpu host-model-only \
    --os-variant "${OS_VARIANT_OPN}" \
    --graphics none \
    --noautoconsole \
    --disk path="${disk_opn}",format=qcow2,bus=virtio \
    --cdrom "${OPN_ISO}" \
    --network network=default,model=virtio \
    --network network="${NET_HARMONYLAN}",model=virtio \
    --boot uefi

  echo "OPNsense VM created. Connect with: sudo virsh console ${VM_OPN}"
  echo "In OPNsense, assign WAN to the NIC with DHCP (default), and LAN to the ${NET_HARMONYLAN} NIC."
  echo "Set LAN IP to ${VLAN_GW}/24 and enable DHCP on LAN (e.g., ${VLAN_GW%.*}.100 - ${VLAN_GW%.*}.200)."

  # --- Create PXE Client VM ---
  local disk_pxe="${IMG_DIR}/${VM_PXE}.qcow2"
  if [[ ! -f "$disk_pxe" ]]; then
    qemu-img create -f qcow2 "$disk_pxe" "${DISK_PXE_GB}G"
  fi

  echo "Creating PXE client VM..."
  virt-install \
    --connect "${CONNECT_URI}" \
    --name "${VM_PXE}" \
    --ram "${RAM_PXE}" \
    --vcpus "${VCPUS_PXE}" \
    --cpu host-model-only \
    --os-variant "${OS_VARIANT_LINUX}" \
    --graphics none \
    --noautoconsole \
    --disk path="${disk_pxe}",format=qcow2,bus=virtio \
    --network network="${NET_HARMONYLAN}",model=virtio \
    --pxe \
    --boot uefi,menu=on

  echo "PXE VM created. It will attempt to netboot on ${NET_HARMONYLAN}."
}

# --- Script Entrypoint ---
case "${1:-}" in
  up)
    mkdir -p "${IMG_DIR}" "${STATE_DIR}"
    download_if_missing "$OPN_URL" "$OPN_ISO"
    download_if_missing "$CENTOS_URL" "$CENTOS_ISO"
    create_lab_environment
    echo "Lab setup complete. Use 'sudo virsh list --all' to see VMs."
    ;;
  clean)
    destroy_vm "${VM_PXE}"
    destroy_vm "${VM_OPN}"
    destroy_network "${NET_HARMONYLAN}"
    # Optionally destroy the default network if you want a full reset
    # destroy_network "default"
    echo "Cleanup complete."
    ;;
  *)
    echo "Usage: sudo $0 {up|clean}"
    exit 1
    ;;
esac