#!/usr/bin/env 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 bind9p /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 bhyverc create-disk zdata/vm/poudriere /vm/poudriere 10
#   doas bhyverc start poudriere zdata/vm/poudriere /vm/poudriere /vm/iso/FreeBSD-13.2-RELEASE-amd64-bootonly.iso
#   doas bhyverc start poudriere zdata/vm/poudriere /vm/poudriere


: ${VERBOSE:="NO"} # or YES
if [ "$VERBOSE" = "YES" ]; then
    set -x
fi

: ${CPU_CORES:="1"}
: ${MEMORY:="1G"}
: ${NETWORK:="NAT"} # or RAW or BOTH
: ${IP_RANGE:="10.215.1.1/24"} # Ignored for RAW networks
: ${INTERFACE_NAME:="jail_nat"} # or the external interface like lagg0 for RAW networks
: ${BRIDGE_NAME:="bridge_$INTERFACE_NAME"} # or bridge_raw for RAW networks
: ${VNC_ENABLE:="NO"}
: ${VNC_LISTEN:="127.0.0.1:5900"}
: ${VNC_WIDTH:="1920"}
: ${VNC_HEIGHT:="1080"}
: ${BIND9P:=""}
: "${CD:=}"

: ${SHUTDOWN_TIMEOUT:="600"} # 10 minutes



############## Setup #########################


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

function log {
    (>&2 echo "${@}")
}

############## Program #########################

function main {
    local cmd
    cmd=$1
    shift
    if [ "$cmd" = "start" ]; then
        init
        start "${@}"
    elif [ "$cmd" = "stop" ]; then
        init
        stop "${@}"
    elif [ "$cmd" = "status" ]; then
        init
        status "${@}"
    elif [ "$cmd" = "console" ]; then
        init
        console "${@}"
    elif [ "$cmd" = "_start_body" ]; then
        init
        start_body "${@}"
    elif [ "$cmd" = "create-disk" ]; then
        create_disk "${@}"
    else
        (>&2 echo "Unknown command: $cmd")
        exit 1
    fi
}

function start {
    local num_vms="$#"
    if [ "$num_vms" -eq 0 ]; then
        log "No VMs specified."
        return 0
    fi

    while [ "$#" -gt 0 ]; do
        local name="$1"
        shift 1
        log "Starting VM $name."
        start_one "$name"
        [ "$#" -eq 0 ] || sleep 5
    done
}

function start_one {
    local name="$1"
    local tmux_name="$name"
    /usr/local/bin/tmux new-session -d -s "$tmux_name" "$0" "_start_body" "$name"
    # /usr/local/bin/tmux new-session -d -s "$tmux_name" "/usr/bin/env VNC_ENABLE=NO VNC_LISTEN=0.0.0.0:5900 /usr/local/bin/bash /home/talexander/launch_opnsense.bash"
}

function launch_pidfile {
    local pidfile="$1"
    shift 1
    mkdir -p "$(dirname "$pidfile")"
    cat > "${pidfile}" <<< "$$"
    set -x
    exec "${@}"
}
export -f launch_pidfile

function stop {
    local num_vms="$#"
    if [ "$num_vms" -eq 0 ]; then
        log "No VMs specified."
        return 0
    fi

    while [ "$#" -gt 0 ]; do
        local name="$1"
        shift 1
        log "Stopping VM $name."
        stop_one "$name"
        [ "$#" -eq 0 ] || sleep 5
    done
}

function stop_one {
    local name="$1"
    local pidfile="/run/bhyverc/${name}/pid"

    if [ ! -e "$pidfile" ]; then
        log "Pid file $pidfile does not exist."
        return 0
    fi

    local bhyve_pid
    bhyve_pid=$(cat "$pidfile")

    if ps -p "$bhyve_pid" >/dev/null; then
        # Send ACPI shutdown command
        log "Sending ACPI shutdown to ${name}:${bhyve_pid}."
        kill -SIGTERM "$bhyve_pid"
    fi

    local timeout_start timeout_end
    timeout_start=$(date +%s)
    while ps -p "$bhyve_pid" >/dev/null; do
        timeout_end=$(date +%s)
        if [ $((timeout_end-timeout_start)) -ge "$SHUTDOWN_TIMEOUT" ]; then
            log "${name}:${bhyve_pid} took more than $SHUTDOWN_TIMEOUT seconds to shut down. Hard powering down."
            break
        fi

        log "Waiting for ${name}:${bhyve_pid} to exit."
        sleep 2
    done

    bhyvectl "--vm=$name" --destroy || true

    local timeout_start timeout_end
    timeout_start=$(date +%s)
    while ps -p "$bhyve_pid" >/dev/null; do
        timeout_end=$(date +%s)
        if [ $((timeout_end-timeout_start)) -ge "$SHUTDOWN_TIMEOUT" ]; then
            log "${name}:${bhyve_pid} took more than $SHUTDOWN_TIMEOUT seconds to hard power down. Giving up."
            break
        fi

        log "Waiting for ${name}:${bhyve_pid} to hard power down."
        sleep 2
    done

    rm -f "$pidfile"

    log "Finished stopping $name."
}

function status {
    local num_vms="$#"

    if [ "$num_vms" -gt 0 ]; then
        for name in "$@"; do
            status_one "$name"
        done
    else
        log "No VMs specified."
    fi
}

function status_one {
    local name="$1"
    local pidfile="/run/bhyverc/${name}/pid"

    if [ ! -e "$pidfile" ]; then
        log "$name is not running."
        return 0
    fi

    local bhyve_pid
    bhyve_pid=$(cat "$pidfile")

    if ! ps -p "$bhyve_pid" >/dev/null; then
        log "$name is not running."
        return 0
    fi

    log "$name is running as pid $bhyve_pid."
}

function console {
    local num_vms="$#"

    if [ "$num_vms" -gt 0 ]; then
        for name in "$@"; do
            log "Attaching to console of VM $name."
            console_one "$name"
        done
    else
        log "No VMs specified."
    fi
}

function console_one {
    local name="$1"
    local tmux_name="$name"
    exec tmux a -t "$tmux_name"
}

function init {
    mkdir -p /run/bhyverc
}

############## Bhyve ###########################

function create_disk {
    local zfs_path="$1"
    local mount_path="$2"
    local 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="$CPU_CORES"
MEMORY="$MEMORY"
NETWORK="$NETWORK"
IP_RANGE="$IP_RANGE"
BRIDGE_NAME="$BRIDGE_NAME"
INTERFACE_NAME="$INTERFACE_NAME"
EOF
    zfs create -s "-V${gigabytes}G" -o volmode=dev -o primarycache=metadata -o secondarycache=none "$zfs_path/disk0"
}

function start_body {
    local name="$1"
    local zfs_path="zdata/vm/$name"
    local mount_path="/vm/$name"
    local mount_cd="$CD"

    if [ -e "${mount_path}/settings" ]; then
        source "${mount_path}/settings"
    fi

    local host_interface_name="$INTERFACE_NAME" # for raw, external interface
    local bridge_name="$BRIDGE_NAME"
    local ip_range="$IP_RANGE" # for raw this value does not matter

    local mac_address
    mac_address=$(calculate_mac_address "$name")

    local additional_args=()

    if [ "$NETWORK" = "NAT" ]; then
        assert_bridge "$host_interface_name" "$bridge_name" "$ip_range"
        local bridge_link_name=$(detect_available_link "${bridge_name}")
        additional_args+=("-s" "2:0,virtio-net,netgraph,path=${bridge_name}:,peerhook=${bridge_link_name},mac=${mac_address}")
    elif [ "$NETWORK" = "RAW" ]; then
        assert_raw "$host_interface_name" "$bridge_name"
        local bridge_link_name=$(detect_available_link "${bridge_name}")
        additional_args+=("-s" "2:0,virtio-net,netgraph,path=${bridge_name}:,peerhook=${bridge_link_name},mac=${mac_address}")
    elif [ "$NETWORK" = "BOTH" ]; then
        assert_bridge "jail_nat" "$bridge_name" "$ip_range"
        assert_raw "$host_interface_name" "bridge_raw"
        local bridge_link_name=$(detect_available_link "${bridge_name}")
        local raw_bridge_link_name=$(detect_available_link "bridge_raw")
        local raw_mac_address=$(calculate_mac_address "${name}_raw")
        additional_args+=("-s" "2:0,virtio-net,netgraph,path=${bridge_name}:,peerhook=${bridge_link_name},mac=${mac_address}")
        additional_args+=("-s" "3:0,virtio-net,netgraph,path=bridge_raw:,peerhook=${raw_bridge_link_name},mac=${raw_mac_address}")
    else
        die 1 "Unrecognized NETWORK type $NETWORK"
    fi

    if [ -n "$BIND9P" ]; then
        additional_args+=("-s" "28,virtio-9p,bind9p=${BIND9P}")
    fi


    # -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
    if [ -n "$mount_cd" ]; then
        additional_args+=("-s" "5,ahci-cd,$mount_cd")
    fi
    if [ "$VNC_ENABLE" = "YES" ]; then
        additional_args+=("-s" "29,fbuf,tcp=$VNC_LISTEN,w=$VNC_WIDTH,h=$VNC_HEIGHT")
    fi
    vms+=("$name")
    while true; do
        local pidfile="/run/bhyverc/${name}/pid"
        trap "set +e; stop_one '${name}'" EXIT

        local launch_cmd=()
        launch_cmd+=(
            launch_pidfile "$pidfile"
            bhyve
            -D
            -c "$CPU_CORES"
            -m "$MEMORY"
            -S
            -H
            -o 'rtc.use_localtime=false'
            -s "0,hostbridge"
            -s "4,nvme,/dev/zvol/${zfs_path}/disk0"
            -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"
        )
        set +e
        rm -f "$pidfile"
        (
            IFS=$' \n\t'
            set -ex
            bash -c "${launch_cmd[*]}"
        )
        local exit_code=$?
        log "Exit code ${exit_code}"
        set -e
        if [ $exit_code -eq 0 ]; then
            echo "Rebooting."
            sleep 5
        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
}

function detect_available_link {
    local bridge_name="$1"
    local linknum=1
    while true; do
        local 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 {
    local host_interface_name="$1"
    local bridge_name="$2"
    local 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 {
    local extif="$1"
    local bridge_name="$2"

    kldload -n ng_bridge ng_eiface ng_ether

    if ! ng_exists "${bridge_name}:"; then
        ngctlcat <<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 {
    local name="$1"
    local source
    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 {
    local start_port="$1"
    local port="$start_port"
    while true; do
        sockstat -P tcp -p 443
        port=$((port + 1))
    done
}

function ngctlcat {
    if [ "$VERBOSE" = "YES" ]; then
        tee /dev/tty | ngctl -d -f -
    else
        ngctl -d -f -
    fi
}

main "${@}"