#!/bin/bash
set -eu

COMMAND=($@)
LOCKFILE=/var/run/netvm.pid
QGA_SOCK_PATH=/var/run/netvm_qga.sock
QGA_SOCK_FDS=false

qga_connect() {
	[[ -v "QGA_SOCK_FDS_PID" ]] && return 0
	coproc QGA_SOCK_FDS (
		exec socat - UNIX-CONNECT:"${QGA_SOCK_PATH}"
	) || return 1
	QGA_SOCK_IN=${QGA_SOCK_FDS[0]}
	QGA_SOCK_OUT=${QGA_SOCK_FDS[1]}
}

qga_execute() {
	local cmd args
	cmd="$1"
	args="${2-}"
	GA_RETURN=""

	jq \
		-ncM \
		--arg cmd "$cmd" \
		--argjson args "$args" \
		'{"execute": $cmd, "arguments": $args}' \
	>&$QGA_SOCK_OUT

	local LINE
	read -t 5 -r -u $QGA_SOCK_IN LINE || return 1

	local ERROR=$(jq -r '.error.desc // empty' <<< "$LINE")
	if [[ -n "$ERROR" ]]; then
		echo "$ERROR" >&2
		return 1
	fi

	GA_RETURN=$(jq -cM .return <<< "$LINE")
}

qga_flush() {
	#Docs say this should work, but it just get parse errors
	#LC_ALL= LC_CTYPE=en_US.UTF-8 printf '%b' "\uff" >&$QGA_SOCK_OUT
	#read -t 5 -r -u $QGA_SOCK_IN LINE
	until ! read -t 1 -r -u $QGA_SOCK_IN LINE; do sleep 0.1; done
}

qga() {
	local cmd args
	cmd="$1"
	args="$2"

	qga_connect

	local id=$((1 + $RANDOM % 10000000))
	qga_execute guest-sync "$(jq -ncM --argjson id "$id" '{"id": $id}')";
	[[ "$(jq -re . <<< "$GA_RETURN")" = "$id" ]] || (echo "Error: guest-sync mismatch" >&2 && return 1)

	unset GA_RETURN
	qga_execute "$cmd" "$args"
	echo "$GA_RETURN" 2>&1

	local RETURN
	kill -INT "$QGA_SOCK_FDS_PID" 2>/dev/null
	wait "$QGA_SOCK_FDS_PID" || RETURN=$?
	if [[ $RETURN != 130 ]]; then
		return $RETURN
	fi
}

function cmd_start(){
        [ ! -f "${LOCKFILE}" ] || { echo "Error: Netvm already running"; exit 1; }
	local net_args=""
	while read line; do
		[ -n "$line" ] || break
		echo "Preparing PCI network device: $line"
		bus=$(echo $line | awk '{ print $1 }')
		id=$(echo $line | awk '{ print $4 }' | tr ':' ' ')
		echo "$id" > /sys/bus/pci/drivers/vfio-pci/new_id 2>&1 || :
		net_args="$net_args -device vfio-pci,host=${bus}"
	done <<< "$(lspci | grep 0200 )"
	if [[ -n "$net_args" ]]; then
		echo Y > /sys/module/vfio_iommu_type1/parameters/allow_unsafe_interrupts
	fi
        printf "Starting netvm...";
        qemu-system-x86_64 \
		-m 512M \
		--machine q35 \
		-nographic \
		-serial none \
		-monitor none \
		-net none \
		-cdrom /guest.img \
		-boot order=d \
		-chardev socket,path=/var/run/netvm_qga.sock,server=on,wait=off,id=qga0 \
		$net_args \
		-device qemu-xhci \
		-device virtio-serial \
		-device virtserialport,chardev=qga0,name=org.qemu.guest_agent.0 \
	>/dev/null 2>&1 &!
	pid=$!
	printf "done\n"
	echo "$pid" > "${LOCKFILE}"
	printf "QGA Socket starting... "
	until [ -S "${QGA_SOCK_PATH}" ]; do sleep 1; done
	printf "done\n"
	printf "Connecting to QGA socket... "
	until qga_connect; do sleep 1; done
	printf "done\n"
	[ -f "/proc/${pid}/status" ] || {
		echo "Error: netvm exited unexpectedly";
		rm "${LOCKFILE}"
		exit 1;
	}

	local id;
	local spin='-\|/'
	local i=0;
	while true; do
		i=$(( (i+1) %4 ));
		printf "\rConnecting to QGA agent... ${spin:$i:1}"
		qga_execute guest-ping "{}" > /dev/null || continue && break
		sleep 1
	done;
	printf "\rConnecting to guest agent... done\n"
	printf "Flushing buffers..."
	qga_flush
	printf "done\n"
	local i=0;
	while true; do
		i=$(( (i+1) %4 ));
		printf "\rSyncing with guest... ${spin:$i:1}"
		id=$((1 + $RANDOM % 10000000))
		qga_execute guest-sync "$(jq -ncM --argjson id "$id" '{"id": $id}')" || continue
		if [[ "$(jq -re . <<< "$GA_RETURN")" = "$id" ]]; then
			printf "\rSyncing with guest... done\n"
			break
		fi;
		sleep 1
	done;
	echo "NetVM boot complete"
}

function cmd_stop(){
	kill $(cat "${LOCKFILE}")
	rm "${LOCKFILE}"
}

function cmd_status(){
	qga guest-get-host-name "{}" | jq -r '."host-name"'
	pid=$(qga guest-exec '{"path": "uptime", "capture-output": true}' | jq -r '.pid')
	out=$(qga guest-exec-status "$(jq -n --argjson pid "$pid" '{pid: $pid }')" \
		| jq -r '."out-data"' \
		| base64 -d \
	)
	echo $out
}

function cmd_push(){
	local source="${COMMAND[1]}"
	local dest="${COMMAND[2]}"
	fo_request=$(jq -n --arg dest "$dest" '{"path": $dest, "mode": "w" }')
	handle=$(qga guest-file-open "$fo_request")
	bufb64=$(base64 "$source")
	count=$(cat "$source" | wc -c)
	fw_request=$(jq -n \
		--argjson handle $handle \
		--argjson count $count \
		--arg bufb64 "$bufb64" \
		'{handle: $handle, "buf-b64": $bufb64, count: $count }' \
	)
	qga guest-file-write "$fw_request"
	fh_request=$(jq -n --argjson handle $handle '{handle: $handle}' )
	qga guest-file-flush "$fh_request"
	qga guest-file-close "$fh_request"
}

function cmd_pull(){
	local source="${COMMAND[1]}"
	local dest="${COMMAND[2]}"
	fo_request=$(jq -n --arg source "$source" '{"path": $source}')
	handle=$(qga guest-file-open "$fo_request")
	fr_request=$(jq -n \
		--argjson handle $handle \
		'{handle: $handle, count: 48000000 }' \
	)
	out=$(qga guest-file-read "$fr_request")
	echo $out | jq -r '."buf-b64"' | base64 -d > $dest
}

function cmd_run(){
	[ -z "${COMMAND[1]}" ] && { echo "Error: missing command"; exit 1; }
	[ -f "${LOCKFILE}" ] || { echo "Error: Netvm is not running"; exit 1; }
	[ -S "${QGA_SOCK_PATH}" ] || { echo "Error: Netvm QGA socket is missing"; exit 1; }
	local cmd="${COMMAND[1]}"
	local args="${COMMAND[@]:2}"
	local args_json="[]"
	if [[ -n "$args" ]]; then
	        args_json=$(printf '%s\n' "$args" | jq -R . | jq -s .)
	fi
	local request
	request=$( \
		jq -n \
			--arg path "$cmd" \
			--argjson args "$args_json" \
			'{
				path: $path,
				arg: $args,
				"capture-output": true
			}' \
	)
	pid=$(qga guest-exec "$request" | jq -r '.pid')
	local exited=false
	until [ "$exited" == "true" ]; do \
		out=$(qga guest-exec-status "$(jq -n --argjson pid "$pid" '{pid: $pid }')" )
		exited=$(echo $out | jq -r '.exited')
		if $exited && jq -r 'has("out-data")' >/dev/null < <(echo $out); then
			echo "$out" | jq -r '."out-data"' | base64 -d
			break
		fi
		sleep 1
	done
}

cmd_usage() {
	cat <<-_EOF
	netvm

	Control network vm headlessly via QMP protocol

	Usage:
	  netvm start
	      Start headless network vm in the background
	  netvm stop
	      Stop headless network vm
	  netvm status
	      Get hostname and uptime from running network vm
	  netvm push <local-path> <remote-path>
	      Push a local file to the network VM
	  netvm pull <remote-path> <local-path>
	      Pull a file from the network VM
	  netvm run "<command>"
	      Run a command in network vm and get stdout
	_EOF
}

case "$1" in
	status) shift; cmd_status $@ ;;
	start)  shift; cmd_start  $@ ;;
	stop)   shift; cmd_stop   $@ ;;
	push)   shift; cmd_push   $@ ;;
	pull)   shift; cmd_pull   $@ ;;
	run)    shift; cmd_run    $@ ;;
	help)   shift; cmd_usage  $@ ;;
	*)             cmd_usage  $@ ;;
esac