#!/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

# Example usage:
#
#   doas qemurc create-disk mint 10
#   doas env CD=/vm/iso/linuxmint-22.2-cinnamon-64bit.iso qemurc start mint
#   doas qemurc start mint
#   doas env WAYLAND_DISPLAY="$XDG_RUNTIME_DIR/$WAYLAND_DISPLAY" XDG_RUNTIME_DIR=/run/user/0 qemurc start mint


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

: ${CPU_CORES:="1"}
: ${MEMORY:="1G"}
: ${GTK_ENABLE:="NO"} # Only enable one, either GTK or VNC
: ${VNC_ENABLE:="NO"} # Only enable one, either GTK or VNC
: ${VNC_LISTEN:="127.0.0.1:0"}
: ${VNC_WIDTH:="1920"}
: ${VNC_HEIGHT:="1080"}
: ${AUDIO_ENABLE:="NO"}
: ${TPM_ENABLE:="NO"}
: ${BIND9P:=""}
: "${CD:=}"

: ${SHUTDOWN_TIMEOUT:="600"}
: ${MOUNT_ROOT:="@mount_root@"}
: ${ZFS_ROOT:="@zfs_root@"}



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


function cleanup {
    sync

    for p in "${pids[@]}"; do
        log "Killing $p"
        kill "$p"
        log "Killed $p"
    done

    for vm in "${vms[@]}"; do
        log "Stopping $vm"
        stop_one "$vm"
        log "Stopped $vm"
    done
}
pids=()
vms=()
trap "set +e; cleanup" EXIT

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"
    tmux new-session -d -s "$tmux_name" "$0" "_start_body" "$name"
}

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/qemurc/${name}/pid"

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

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

    if ps -p "$qemu_pid" >/dev/null; then
        # We cannot send a graceful shutdown command externally to qemu: https://gitlab.com/qemu-project/qemu/-/issues/148
        log "Killing ${name}:${qemu_pid}."
        kill -SIGTERM "$qemu_pid"
    fi

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

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

    kill -9 "$qemu_pid"

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

        log "Waiting for ${name}:${qemu_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/qemurc/${name}/pid"

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

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

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

    log "$name is running as pid $qemu_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/qemurc
}

############## qemu ############################

function create_disk {
    local name="$1"
    local gigabytes="$2"

    local zfs_path="${ZFS_ROOT}/${name}"
    local mount_path="${MOUNT_ROOT}/${name}"

    zfs create -o mountpoint=none -o canmount=off "$zfs_path"
    zfs create -o "mountpoint=$mount_path" -o canmount=on "$zfs_path/settings"
    zfs create -s "-V${gigabytes}G" -o volmode=dev -o primarycache=metadata -o secondarycache=none "$zfs_path/disk0"
    zfs snapshot -r "$zfs_path@empty"

    install -m0600 "@OVMFfd@/FV/OVMF_VARS.fd" "${mount_path}/"
    tee "${mount_path}/settings" <<EOF
CPU_CORES="$CPU_CORES"
MEMORY="$MEMORY"
GTK_ENABLE="$GTK_ENABLE"
VNC_ENABLE="$VNC_ENABLE"
VNC_LISTEN="$VNC_LISTEN"
VNC_WIDTH="$VNC_WIDTH"
VNC_HEIGHT="$VNC_HEIGHT"
AUDIO_ENABLE="$AUDIO_ENABLE"
TPM_ENABLE="$TPM_ENABLE"
BIND9P="$BIND9P"
EOF

}

function start_body {
    local name="$1"
    local zfs_path="${ZFS_ROOT}/${name}"
    local mount_path="${MOUNT_ROOT}/${name}"
    local run_path="/run/qemurc/${name}"
    local mount_cd="$CD"
    local swtpm_sock="${run_path}/swtpm.sock"
    local swtpm_path="${MOUNT_ROOT}/${name}/swtpm"

    install -d -m 0700 "$run_path"

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

    local additional_args=()

    if [ -n "$BIND9P" ]; then
        additional_args+=(-device "virtio-9p-type,fsdev=${BIND9P},mount_tag=bind9p")
    fi

    if [ -n "$mount_cd" ]; then
        additional_args+=(-cdrom "$mount_cd")
    fi
    if [ "$VNC_ENABLE" = "YES" ]; then
        additional_args+=(-vnc "${VNC_LISTEN},power-control=on")
    fi

    if [ "$AUDIO_ENABLE" = "YES" ]; then
        additional_args+=(-audio "driver=pa,model=virtio,server=/run/user/11235/pulse/native")
    fi

    if [ "$TPM_ENABLE" = "YES" ]; then
        install -d -m 0700 "$swtpm_path"
        swtpm socket --tpm2 --tpmstate dir="$swtpm_path" --ctrl type=unixio,path="$swtpm_sock" &
        local tpm_pid=$!
        pids+=("$tpm_pid")
        additional_args+=(-chardev "socket,id=chrtpm,path=$swtpm_sock"
                          -tpmdev "emulator,id=tpm0,chardev=chrtpm"
                          -device "tpm-tis,tpmdev=tpm0")
    fi

    if [ "$GTK_ENABLE" = "YES" ]; then
        additional_args+=(
            -device 'virtio-gpu-gl,hostmem=8G,blob=true,venus=true'
            -display 'gtk,gl=on'
            -vga virtio
        )
    fi


    vms+=("$name")

    local pidfile="/run/qemurc/${name}/pid"

    local launch_cmd=()
    launch_cmd+=(
        launch_pidfile "$pidfile"
        qemu-system-x86_64
        -accel kvm
        -cpu host
        -smp cores="$CPU_CORES"
        -m "$MEMORY"
        -rtc base=localtime
        -drive "file=\"@OVMFfd@/FV/OVMF_CODE.fd\",if=pflash,format=raw,readonly=on"
        -drive "if=pflash,format=raw,file=\"$(readlink -f "${mount_path}/OVMF_VARS.fd")\""
        -drive "if=none,file=/dev/zvol/${zfs_path}/disk0,format=raw,id=hd0"
        -device 'nvme,serial=deadbeef,drive=hd0'
        -nic 'user,hostfwd=tcp::60022-:22'
        -boot order=d
        "${additional_args[@]}"
    )
    set +e
    rm -f "$pidfile"
    (
        IFS=$' \n\t'
        set -ex
        bash -c "${launch_cmd[*]}"
    )
    local exit_code=$?
    log "Exit code ${exit_code}"
    set -e
}

main "${@}"