#!/usr/local/bin/bash
#
set -euo pipefail
IFS=$'\n\t'
DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"

# Share a host directory to the guest via 9pfs.
#
# Inside the VM run:
#   mount -t virtfs -o trans=virtio sharename /some/vm/path
#   mount -t 9p -o cache=mmap -o msize=512000 sharename /mnt/9p
#   mount -t 9p -o trans=virtio,cache=mmap,msize=512000 sharename /path/to/mountpoint
# bhyve_options="-s 28,virtio-9p,sharename=/"

# Enable Sound
# bhyve_options="-s 16,hda,play=/dev/dsp,rec=/dev/dsp"

# Example usage:
#
#   doas bhyve_netgraph_bridge create-disk zdata/vm/poudriere /vm/poudriere 10
#   doas bhyve_netgraph_bridge start poudriere zdata/vm/poudriere /vm/poudriere jail_nat 10.215.1.1/24 /vm/iso/FreeBSD-13.2-RELEASE-amd64-bootonly.iso
#   doas bhyve_netgraph_bridge start poudriere zdata/vm/poudriere /vm/poudriere jail_nat 10.215.1.1/24

function main {
    if [ "$1" = "create-disk" ]; then
        shift 1
        create_disk "${@}"
    elif [ "$1" = "start" ]; then
        shift 1
        start_vm "${@}"
    else
        >&2 echo "Unrecognized command"
        exit 1
    fi
}

function die {
    local status_code="$1"
    shift
    (>&2 echo "${@}")
    exit "$status_code"
}

function create_disk {
    zfs_path="$1"
    mount_path="$2"
    gigabytes="$3"
    zfs create -o "mountpoint=$mount_path" "$zfs_path"
    cp /usr/local/share/edk2-bhyve/BHYVE_UEFI_VARS.fd "${mount_path}/"
    tee "${mount_path}/settings" <<EOF
CPU_CORES=1
MEMORY=1G
NETWORK=NAT
EOF
    zfs create -s "-V${gigabytes}G" -o volmode=dev "$zfs_path/disk0"
}

function start_vm {
    name="$1"
    zfs_path="$2"
    mount_path="$3"
    host_interface_name="$4" # for raw, external interface
    bridge_name="bridge_${host_interface_name}"
    ip_range="$5" # for raw this value does not matter
    mount_cd="${6:-}"

    mac_address=$(calculate_mac_address "$name")

    CPU_CORES=1
    MEMORY=1G
    NETWORK="NAT"
    if [ -e "${mount_path}/settings" ]; then
        source "${mount_path}/settings"
    fi

    if [ "$NETWORK" = "NAT" ]; then
        assert_bridge "$host_interface_name" "$bridge_name" "$ip_range"
    elif [ "$NETWORK" = "RAW" ]; then
        assert_raw "$host_interface_name" "$bridge_name"
    else
        die 1 "Unrecognized NETWORK type $NETWORK"
    fi
    bridge_link_name=$(detect_available_link "${bridge_name}")


    # -H release the CPU when guest issues HLT instruction. Otherwise 100% of core will be consumed.
         # -s 3,ahci-cd,/vm/.iso/archlinux-2023.04.01-x86_64.iso \
             # -s 29,fbuf,tcp=0.0.0.0:5900,w=1920,h=1080,wait \
            # -s 29,fbuf,tcp=0.0.0.0:5900,w=1920,h=1080 \

    # TODO: Look into using nmdm instead of stdio for serial console
    additional_args=()
    if [ -n "$mount_cd" ]; then
        additional_args+=("-s" "3,ahci-cd,$mount_cd")
    fi
    while true; do
        set -x
        set +e
        bhyve \
            -D \
            -c $CPU_CORES \
            -m $MEMORY \
            -H \
            -s 0,hostbridge \
            -s "4,nvme,/dev/zvol/${zfs_path}/disk0" \
            -s "2:0,virtio-net,netgraph,path=${bridge_name}:,peerhook=${bridge_link_name},mac=${mac_address}" \
            -s 30,xhci,tablet \
            -s 31,lpc -l com1,stdio \
            -l "bootrom,/usr/local/share/uefi-firmware/BHYVE_UEFI.fd,${mount_path}/BHYVE_UEFI_VARS.fd" \
            "${additional_args[@]}" \
            "$name"
        exit_code=$?
        set -e
        set +x
        if [ $exit_code -eq 0 ]; then
            echo "Rebooting."
        elif [ $exit_code -eq 1 ]; then
            echo "Powered off."
            break
        elif [ $exit_code -eq 2 ]; then
            echo "Halted."
            break
        elif [ $exit_code -eq 3 ]; then
            echo "Triple fault."
            break
        elif [ $exit_code -eq 4 ]; then
            echo "Exited due to an error."
            break
        fi
    done

    bhyvectl "--vm=$name" --destroy
    echo "Destroyed bhyve vm."
}

function detect_available_link {
    bridge_name="$1"
    linknum=1
    while true; do
        link_name="link${linknum}"
        if ! ng_exists "${bridge_name}:${link_name}"; then
            echo "$link_name"
            return
        fi
        linknum=$((linknum + 1))
        if [ "$linknum" -gt 90 ]; then
            (>&2 echo "No available links on bridge $bridge_name")
            exit 1
        fi
    done
}

function assert_bridge {
    host_interface_name="$1"
    bridge_name="$2"
    ip_range="$3"

    if ! ng_exists "${bridge_name}:"; then
        ngctl -d -f - <<EOF
mkpeer . eiface hook ether
name .:hook $host_interface_name
EOF
        ngctl -d -f - <<EOF
mkpeer ${host_interface_name}: bridge ether link0
name ${host_interface_name}:ether $bridge_name
EOF
        ifconfig $(ngctl msg "${host_interface_name}:" getifname | grep Args | cut -d '"' -f 2) name "${host_interface_name}" "$ip_range" up
    fi
}

function assert_raw {
    extif="$1"
    bridge_name="$2"

    kldload -n ng_bridge ng_eiface ng_ether

    if ! ng_exists "${bridge_name}:"; then
        ngctl -d -f - <<EOF
# Create a bridge.
mkpeer $extif: bridge lower link0
# Assign a name to the bridge.
name $extif:lower ${bridge_name}
# Since the host is also using $extif, we need to connect the upper hook also. Otherwise we will lose connectivity.
connect $extif: ${bridge_name}: upper link1

# Enable promiscuous mode so the host ethernet adapter accepts packets for all addresses
msg $extif: setpromisc 1

# Do not overwrite source address on packets
msg $extif: setautosrc 0
EOF
    fi
}

function ng_exists {
    ngctl status "${1}" >/dev/null 2>&1
}

function calculate_mac_address {
    name="$1"
    source=$(md5 -r -s "$name" | awk '{print $1}')
    echo "06:${source:0:2}:${source:2:2}:${source:4:2}:${source:6:2}:${source:8:2}"
}

function find_available_port {
    start_port="$1"
    port="$start_port"
    while true; do
        sockstat -P tcp -p 443
        port=$((port + 1))
    done
}


main "${@}"