#!/usr/bin/env bash
# vjunos-evo-lab.sh - Spin up a 3-router vJunosEvolved lab on KVM
#
# Usage: sudo ./vjunos-evo-lab.sh {setup|start|stop|destroy|status|console N}
#
# vJunosEvolved is a single (non-nested) VM running Junos OS Evolved as a
# virtualized PTX10001-36MR. Works fine on AMD because there's no inner KVM.
#
# First time:
#   sudo ./vjunos-evo-lab.sh setup    (one-time: bridges, disks, configs)
#   sudo ./vjunos-evo-lab.sh start    (boot all 3 routers)
#   sudo ./vjunos-evo-lab.sh console 1   (connect to rt1)
#
# Subsequent times: just 'start' / 'stop'.
#
# Topology: triangle
#   rt1:et-0/0/0  <->  rt2:et-0/0/0   (vjlink-12)
#   rt2:et-0/0/1  <->  rt3:et-0/0/0   (vjlink-23)
#   rt1:et-0/0/1  <->  rt3:et-0/0/1   (vjlink-13)
#   All re0:mgmt-0 share vjevo-mgmt bridge.

set -euo pipefail

# =================== CONFIG ===================
LAB_DIR="${LAB_DIR:-/var/lib/libvirt/images/vjunos-evo-lab}"
BASE_IMAGE="${BASE_IMAGE:-/home/jack/Downloads/ISO/Juniper/vJunosEvolved-25.4R1.13-EVO.qcow2}"
NUM_ROUTERS=3
RAM_MB=8192             # vJunosEvolved minimum
VCPUS=4
MGMT_BRIDGE="vjevo-mgmt"
CONSOLE_BASE_PORT=8700  # different from vjunos-switch ports to avoid collisions
ROOT_PASSWORD="Juniper123"

# Inter-router links: "rtA:portA rtB:portB bridge-name"
# port numbers are 0-indexed and map to et-0/0/<n>
LINKS=(
    "1:0 2:0 vjlink-12"
    "2:1 3:0 vjlink-23"
    "1:1 3:1 vjlink-13"
)
# ==============================================

usage() {
    cat <<EOF
Usage: sudo $0 {setup|start|stop|destroy|status|console N}

  setup       Create bridges, copy disks, generate configs, define VMs
  start       Start all routers (idempotent)
  stop        Gracefully shut down all routers
  destroy     Stop, undefine, and delete ALL lab files (asks for confirm)
  status      Show running state, console ports, topology
  console N   Telnet to router N's console (Ctrl-] then 'quit' to exit)

Base image: $BASE_IMAGE
Lab dir:    $LAB_DIR
EOF
}

log() { echo "[$(date +%H:%M:%S)] $*"; }
err() { echo "ERROR: $*" >&2; exit 1; }

require_root() { [[ $EUID -eq 0 ]] || err "needs root: sudo $0 $*"; }

require_tools() {
    for t in virsh qemu-img mkfs.vfat mcopy mmd ip telnet; do
        command -v "$t" >/dev/null || err "missing tool: $t"
    done
    # vJunosEvolved needs UEFI firmware (OVMF) - it has a GPT/EFI partition layout
    if ! ls /usr/share/edk2/ovmf/*.fd /usr/share/OVMF/*.fd 2>/dev/null | head -1 >/dev/null; then
        err "OVMF firmware not found. Install with: sudo dnf install edk2-ovmf"
    fi
}

detect_libvirt_user() {
    if id libvirt-qemu &>/dev/null; then
        LIBVIRT_USER="libvirt-qemu"; LIBVIRT_GROUP="kvm"
    elif id qemu &>/dev/null; then
        LIBVIRT_USER="qemu"; LIBVIRT_GROUP="qemu"
    else
        LIBVIRT_USER="root"; LIBVIRT_GROUP="root"
    fi
}

# --------------- BRIDGES ---------------
create_bridge() {
    local name=$1
    if ! ip link show "$name" &>/dev/null; then
        ip link add "$name" type bridge
        log "  created bridge: $name"
    fi
    ip link set "$name" up
    echo 0 > "/sys/class/net/$name/bridge/stp_state" 2>/dev/null || true
    echo 0x4000 > "/sys/class/net/$name/bridge/group_fwd_mask" 2>/dev/null || true
}

delete_bridge() {
    local name=$1
    if ip link show "$name" &>/dev/null; then
        ip link set "$name" down
        ip link delete "$name"
        log "  deleted bridge: $name"
    fi
}

# --------------- CONFIG DRIVE ---------------
# vJunosEvolved reads /config/juniper.conf from a FAT-formatted USB disk (raw img).
make_config_drive() {
    local rt_num=$1
    local outfile=$2
    local hostname="rt${rt_num}"
    local loopback_ip="192.168.255.${rt_num}"

    local tmpconf
    tmpconf=$(mktemp)
    cat > "$tmpconf" <<EOF
system {
    host-name ${hostname};
    root-authentication {
        plain-text-password-value "${ROOT_PASSWORD}";
    }
    login {
        user admin {
            class super-user;
            authentication {
                plain-text-password-value "${ROOT_PASSWORD}";
            }
        }
    }
    services {
        ssh {
            root-login allow;
        }
        netconf {
            ssh;
        }
    }
    syslog {
        user * { any emergency; }
        file messages { any notice; authorization info; }
        file interactive-commands { interactive-commands any; }
    }
}
interfaces {
    re0:mgmt-0 {
        unit 0 {
            family inet {
                dhcp;
            }
        }
    }
    lo0 {
        unit 0 {
            family inet {
                address ${loopback_ip}/32;
            }
        }
    }
}
EOF

    # vJunosEvolved expects raw format (not qcow2) for the config disk
    truncate -s 16M "$outfile"
    mkfs.vfat -n vmm-data "$outfile" >/dev/null
    mmd -i "$outfile" ::/config
    mcopy -i "$outfile" "$tmpconf" ::/config/juniper.conf
    rm -f "$tmpconf"
    log "  config drive: $outfile (hostname=$hostname, lo0=$loopback_ip/32, root pw=$ROOT_PASSWORD)"
}

# --------------- LIBVIRT XML ---------------
generate_xml() {
    local rt_num=$1
    local console_port=$((CONSOLE_BASE_PORT + rt_num))
    local disk="${LAB_DIR}/rt${rt_num}-disk.qcow2"
    local config="${LAB_DIR}/rt${rt_num}-config.img"
    local xml="${LAB_DIR}/rt${rt_num}.xml"

    declare -A port_to_bridge=()
    for link in "${LINKS[@]}"; do
        read -r endpointA endpointB bridge <<< "$link"
        local rtA="${endpointA%:*}" portA="${endpointA#*:}"
        local rtB="${endpointB%:*}" portB="${endpointB#*:}"
        [[ "$rtA" == "$rt_num" ]] && port_to_bridge[$portA]=$bridge
        [[ "$rtB" == "$rt_num" ]] && port_to_bridge[$portB]=$bridge
    done

    # First NIC = re0:mgmt-0, then et-0/0/0, et-0/0/1, ... in port order
    local interfaces=""
    interfaces+="
    <interface type='bridge'>
      <source bridge='${MGMT_BRIDGE}'/>
      <mac address='52:54:00:0e:00:0${rt_num}'/>
      <model type='virtio'/>
    </interface>"

    for port in $(printf "%s\n" "${!port_to_bridge[@]}" | sort -n); do
        local br="${port_to_bridge[$port]}"
        interfaces+="
    <interface type='bridge'>
      <source bridge='${br}'/>
      <mac address='52:54:00:7e:0${rt_num}:0${port}'/>
      <model type='virtio'/>
    </interface>"
    done

    cat > "$xml" <<EOF
<domain type='kvm'>
  <name>vjunos-rt${rt_num}</name>
  <memory unit='MiB'>${RAM_MB}</memory>
  <currentMemory unit='MiB'>${RAM_MB}</currentMemory>
  <vcpu>${VCPUS}</vcpu>
  <cpu mode='host-passthrough' check='none'>
    <topology sockets='1' cores='${VCPUS}' threads='1'/>
  </cpu>
  <os firmware='efi'>
    <type arch='x86_64' machine='pc'>hvm</type>
    <loader secure='no'/>
    <boot dev='hd'/>
  </os>
  <features>
    <acpi/>
    <apic/>
  </features>
  <clock offset='utc'/>
  <on_poweroff>destroy</on_poweroff>
  <on_reboot>destroy</on_reboot>
  <on_crash>preserve</on_crash>
  <devices>
    <emulator>/usr/bin/qemu-system-x86_64</emulator>
    <video>
      <model type='vga' vram='16384' heads='1' primary='yes'/>
    </video>
    <graphics type='vnc' port='-1' autoport='yes' listen='127.0.0.1'/>
    <disk type='file' device='disk'>
      <driver name='qemu' type='qcow2' cache='writeback'/>
      <source file='${disk}'/>
      <target dev='vda' bus='virtio'/>
    </disk>
    <disk type='file' device='disk'>
      <driver name='qemu' type='raw' cache='writeback'/>
      <source file='${config}'/>
      <target dev='sda' bus='usb'/>
    </disk>${interfaces}
    <serial type='tcp'>
      <source mode='bind' host='127.0.0.1' service='${console_port}'/>
      <protocol type='telnet'/>
      <target port='0'/>
    </serial>
    <console type='tcp'>
      <source mode='bind' host='127.0.0.1' service='${console_port}'/>
      <protocol type='telnet'/>
      <target type='serial' port='0'/>
    </console>
  </devices>
</domain>
EOF
    log "  XML: $xml (console port: $console_port)"
}

# --------------- SUBCOMMANDS ---------------
do_setup() {
    require_root setup
    require_tools
    detect_libvirt_user
    [[ -f "$BASE_IMAGE" ]] || err "base image not found: $BASE_IMAGE"

    log "Setting up lab in $LAB_DIR"
    mkdir -p "$LAB_DIR"

    log "Bridges:"
    create_bridge "$MGMT_BRIDGE"
    for link in "${LINKS[@]}"; do
        read -r _ _ bridge <<< "$link"
        create_bridge "$bridge"
    done

    for i in $(seq 1 $NUM_ROUTERS); do
        log "Router $i:"
        local disk="${LAB_DIR}/rt${i}-disk.qcow2"
        local config="${LAB_DIR}/rt${i}-config.img"

        if [[ ! -f "$disk" ]]; then
            cp --reflink=auto "$BASE_IMAGE" "$disk"
            log "  disk: $disk"
        else
            log "  disk: $disk (exists, skipping)"
        fi

        # Always regenerate config drive in case juniper.conf changed
        make_config_drive "$i" "$config"

        chown "$LIBVIRT_USER:$LIBVIRT_GROUP" "$disk" "$config" 2>/dev/null || true
        chmod 644 "$disk" "$config"

        generate_xml "$i"

        # Undefine first if exists, so XML changes apply
        if virsh dominfo "vjunos-rt${i}" &>/dev/null; then
            local cur_state; cur_state=$(virsh domstate "vjunos-rt${i}")
            if [[ "$cur_state" == "running" ]]; then
                log "  vjunos-rt${i} is running - shutting down before redefine"
                virsh shutdown "vjunos-rt${i}" >/dev/null
                for _ in $(seq 1 60); do
                    [[ "$(virsh domstate vjunos-rt${i} 2>/dev/null)" != "running" ]] && break
                    sleep 1
                done
            fi
            # UEFI domains have NVRAM files - need --nvram to undefine cleanly.
            # Also handle managed-save state if present. Multiple fallbacks for virsh version differences.
            virsh undefine "vjunos-rt${i}" --nvram --managed-save >/dev/null 2>&1 || \
                virsh undefine "vjunos-rt${i}" --nvram >/dev/null 2>&1 || \
                virsh undefine "vjunos-rt${i}" >/dev/null 2>&1 || true
        fi
        virsh define "${LAB_DIR}/rt${i}.xml" >/dev/null
        log "  defined: vjunos-rt${i}"
    done

    log ""
    log "Setup complete. Run 'sudo $0 start' to boot the lab."
}

do_start() {
    require_root start
    for i in $(seq 1 $NUM_ROUTERS); do
        if ! virsh dominfo "vjunos-rt${i}" &>/dev/null; then
            log "vjunos-rt${i} not defined - run setup first"; continue
        fi
        local state; state=$(virsh domstate "vjunos-rt${i}")
        if [[ "$state" == "running" ]]; then
            log "vjunos-rt${i}: already running"
        else
            virsh start "vjunos-rt${i}" >/dev/null
            log "vjunos-rt${i}: started (console: telnet localhost $((CONSOLE_BASE_PORT + i)))"
        fi
    done
    log ""
    log "Routers booting. Junos Evolved takes ~3-5 minutes to reach login prompt."
}

do_stop() {
    require_root stop
    log "Sending graceful shutdown to all routers..."
    for i in $(seq 1 $NUM_ROUTERS); do
        if virsh dominfo "vjunos-rt${i}" &>/dev/null; then
            local state; state=$(virsh domstate "vjunos-rt${i}")
            if [[ "$state" == "running" ]]; then
                virsh shutdown "vjunos-rt${i}" >/dev/null
                log "  vjunos-rt${i}: shutdown sent"
            fi
        fi
    done
    log "Waiting up to 90s for graceful shutdown..."
    for sec in $(seq 1 90); do
        local running=0
        for i in $(seq 1 $NUM_ROUTERS); do
            [[ "$(virsh domstate vjunos-rt${i} 2>/dev/null)" == "running" ]] && running=1
        done
        [[ $running -eq 0 ]] && { log "All stopped after ${sec}s."; return; }
        sleep 1
    done
    log "WARNING: some routers didn't shut down cleanly."
}

do_destroy() {
    require_root destroy
    read -p "This will WIPE all routers and disks. Type YES to continue: " confirm
    [[ "$confirm" == "YES" ]] || { log "aborted"; exit 0; }

    do_stop
    for i in $(seq 1 $NUM_ROUTERS); do
        virsh undefine "vjunos-rt${i}" --nvram --managed-save 2>/dev/null || \
            virsh undefine "vjunos-rt${i}" --nvram 2>/dev/null || \
            virsh undefine "vjunos-rt${i}" 2>/dev/null || true
    done
    rm -f "${LAB_DIR}"/rt*-disk.qcow2 "${LAB_DIR}"/rt*-config.img "${LAB_DIR}"/rt*.xml
    delete_bridge "$MGMT_BRIDGE"
    for link in "${LINKS[@]}"; do
        read -r _ _ bridge <<< "$link"
        delete_bridge "$bridge"
    done
    log "Lab destroyed. Base image preserved: $BASE_IMAGE"
}

do_status() {
    echo "=== Routers ==="
    for i in $(seq 1 $NUM_ROUTERS); do
        if virsh dominfo "vjunos-rt${i}" &>/dev/null; then
            local state; state=$(virsh domstate "vjunos-rt${i}")
            printf "  vjunos-rt%d  %-10s  console: telnet localhost %d   lo0: 192.168.255.%d\n" \
                "$i" "$state" "$((CONSOLE_BASE_PORT + i))" "$i"
        else
            printf "  vjunos-rt%d  %-10s\n" "$i" "not-defined"
        fi
    done
    echo ""
    echo "=== Bridges ==="
    ip -br link show type bridge 2>/dev/null | grep -E "vjevo-mgmt|vjlink-" || echo "  (none)"
    echo ""
    echo "=== Topology ==="
    for link in "${LINKS[@]}"; do
        read -r a b br <<< "$link"
        echo "  rt${a%:*}:et-0/0/${a#*:}  <==[$br]==>  rt${b%:*}:et-0/0/${b#*:}"
    done
}

do_console() {
    local n="${1:-}"
    [[ -z "$n" ]] && err "Usage: $0 console <1|2|3>"
    [[ "$n" =~ ^[0-9]+$ && "$n" -ge 1 && "$n" -le $NUM_ROUTERS ]] || err "router number must be 1-$NUM_ROUTERS"
    local port=$((CONSOLE_BASE_PORT + n))
    log "rt${n} console (port $port). Disconnect: Ctrl-] then 'quit'"
    telnet localhost "$port"
}

# --------------- MAIN ---------------
case "${1:-}" in
    setup)   do_setup ;;
    start)   do_start ;;
    stop)    do_stop ;;
    destroy) do_destroy ;;
    status)  do_status ;;
    console) shift; do_console "$@" ;;
    -h|--help|"") usage ;;
    *)       usage; exit 1 ;;
esac