#!/usr/bin/env bash
# dispatch.sh: turns Amp orb provider jobs into exe.dev VMs, one VM per orb.
# From https://exe.dev/deploy/amp. Based on the exe.dev example in
# https://ampcode.com/docs/orbs/self-hosted, but it calls exe.dev through the
# amp-orbs integration (https://amp-orbs.int.exe.xyz), so no exe.dev key is
# on this VM, and the integration reaches only the VMs it made.
set -euo pipefail

EXE_API="${EXE_API:-https://amp-orbs.int.exe.xyz/exec}"

# Run an exe.dev command through the integration and print its output.
exe() {
	curl -sS --max-time 280 -X POST "$EXE_API" --data-binary "$*"
}

# The orb's VM: running, paused, or empty if there is none.
vms=$(exe ls "$ORB_ID" --json)
if ! jq -e '.vms' <<<"$vms" >/dev/null 2>&1; then
	echo "exe.dev: $vms" >&2
	exit 1
fi
status=$(jq -r --arg id "$ORB_ID" '.vms[] | select(.vm_name == $id) | .status' <<<"$vms")

case "$ORB_OPERATION" in
spawn)
	if [[ -z "$status" ]]; then
		# Orb types you add in Amp arrive as ORB_TYPE; map them to sizes here.
		size=""
		case "${ORB_TYPE-}" in
		large) size="--cpu=4 --memory=16GB" ;;
		esac
		out=$(exe new --name="$ORB_ID" --image="$ORB_IMAGE" $size --no-email --json)
		if jq -e '.error' <<<"$out" >/dev/null; then
			echo "exe.dev new $ORB_ID failed: $out" >&2
			exit 1
		fi
		echo "created VM $ORB_ID"
	elif [[ "$status" == paused ]]; then
		exe resume "$ORB_ID" >/dev/null
		echo "resumed VM $ORB_ID"
	fi

	# Start amp orb bootstrap in the orb as its user, detached, with this
	# job's token. The script goes over base64 so no quoting can break it.
	script=$(
		printf 'export ORB_TOKEN=%q\n' "$ORB_TOKEN"
		if [[ -n "${AMP_URL-}" ]]; then
			printf 'export AMP_URL=%q\n' "$AMP_URL"
		fi
		cat <<'BOOTSTRAP'
sudo -u user --preserve-env=ORB_TOKEN,AMP_URL bash -lc \
    'cd "$HOME"; setsid nohup amp orb bootstrap </dev/null >.amp-bootstrap.log 2>&1 &'
echo bootstrap-started
BOOTSTRAP
	)
	b64=$(base64 -w0 <<<"$script")
	for _ in $(seq 1 30); do
		if exe ssh "$ORB_ID" "echo $b64 | base64 -d | bash" | grep -q bootstrap-started; then
			echo "started amp orb bootstrap in $ORB_ID"
			exit 0
		fi
		sleep 2
	done
	echo "could not reach $ORB_ID to start amp orb bootstrap" >&2
	exit 1
	;;
pause)
	# exe.dev pause keeps memory too. If the integration can't pause (no
	# pause in its --cmds), the orb stays running and the next spawn
	# reconnects it.
	if [[ "$status" == running ]]; then
		out=$(exe pause "$ORB_ID" --json)
		if jq -e '.error' <<<"$out" >/dev/null 2>&1; then
			echo "left $ORB_ID running: $out"
		else
			echo "paused VM $ORB_ID"
		fi
	fi
	;;
delete)
	if [[ -n "$status" ]]; then
		exe rm "$ORB_ID" >/dev/null
		echo "deleted VM $ORB_ID"
	fi
	;;
*)
	echo "unknown ORB_OPERATION: $ORB_OPERATION" >&2
	exit 1
	;;
esac