#!/bin/bash

# shellcheck disable=SC2155
# shellcheck disable=SC2164
readonly script_path="$( cd -- "$(dirname "$0")" >/dev/null 2>&1 ; pwd -P )" # https://stackoverflow.com/a/4774063

# Location of the SSH config to be generated
SSH_CONFIG=${SSH_CONFIG:-"$script_path/ssh.conf"}

NUM_NODES=${NUM_NODES:-"2"}
MEMORY=${MEMORY:-"2048"}
CPU=${CPU:-"2"}
DISK_MAX_SIZE=${DISK_MAX_SIZE:-"20G"}

# Location of VM state information, including its persistent disk
STATE_DIR=${STATE_DIR:-"$script_path/state"}

# Location of Linux cloud image
BASE_IMAGE_FILENAME=${BASE_IMAGE_FILENAME:-"$script_path/images/Rocky-9-GenericCloud-Base.latest.x86_64.qcow2"}

# Architecture that QEMU will emulate
EMU_ARCH=${EMU_ARCH:-"x86_64"}

# Check for prerequisites
if ! command -v mkisofs > /dev/null 2>&1; then
    echo "mkisofs not found."
    exit 1
fi

if ! command -v qemu-system-"$EMU_ARCH" > /dev/null 2>&1; then
    echo "qemu-system-$EMU_ARCH not found."
    exit 1
fi

stop() {
    for ((i=0;i<NUM_NODES;i++))
    do
        pkill -F "$STATE_DIR"/node$i.pid >/dev/null 2>&1
        rm -f "$STATE_DIR"/node$i.pid
    done
}

start() {
    stop

    for ((i=0;i<NUM_NODES;i++))
    do
        if ! [ -f "$STATE_DIR"/node$i.qcow2 ] || ! [ -f "$STATE_DIR"/node$i-cloudinit.iso ]; then
            echo "Missing disk images. Run create first."
            exit 1
        fi

        echo "Starting node $i"
        qemu-system-x86_64 \
            -enable-kvm \
            -cpu host \
            -m "$MEMORY" \
            -smp "$CPU" \
            -device virtio-net-pci,netdev=net$i,mac=52:54:00:00:00:"$(printf '%02x' $i)" \
            -netdev socket,id=net$i,mcast=230.0.0.1:1234 \
            -device virtio-net-pci,netdev=host$i \
            -netdev user,id=host$i,hostfwd=tcp::$((3100 + i))-:22 \
            -drive file="$STATE_DIR"/node$i.qcow2,media=disk,if=virtio \
            -drive file="$STATE_DIR"/node$i-cloudinit.iso,media=cdrom \
            -display none \
            -daemonize \
            -pidfile "$STATE_DIR"/node$i.pid
    done

    for ((i=0;i<NUM_NODES;i++))
    do
        local attempts_remaining=10
        while ! ssh -o ConnectTimeout=2 -o BatchMode=true -F "$SSH_CONFIG" node$i exit >/dev/null 2>&1 && ((attempts_remaining > 0));
        do
            ((attempts_remaining--))
            if [ $attempts_remaining == 0 ]; then
                echo "SSH service on node$i took too long to be ready"
                exit 1
            fi
            echo "Waiting for SSH service on node$i to be ready; $attempts_remaining attempts remaining"
            sleep 5
        done
        echo "SSH service on node$i is up"
    done

    echo "To access a VM, run:"
    echo
    echo "    ssh -F $SSH_CONFIG node#"
    echo
    echo "Where # is the node number you want to reach."
}

destroy() {
    stop

    rm -f "$STATE_DIR"/id_ed25519
    rm -f "$STATE_DIR"/id_ed25519.pub
    rm -f "$SSH_CONFIG"

    for ((i=0;i<NUM_NODES;i++))
    do
        rm -f "$STATE_DIR"/node$i-cloudinit.iso
        rm -f "$STATE_DIR"/node$i.qcow2
    done
}

create() {
    destroy

    mkdir -p "$STATE_DIR"

    ssh-keygen -f qemu_key -N "" -t ed25519 -C "" -f "$STATE_DIR"/id_ed25519 >/dev/null 2>&1
    ssh_pubkey=$(tr -d '\n' < "$STATE_DIR"/id_ed25519.pub)

    for ((i=0;i<NUM_NODES;i++))
    do
        # Create cloud-init volume
        cat > "$STATE_DIR"/user-data <<- EOF
			#cloud-config

			users:
			  - name: root
			    ssh_authorized_keys:
			      - $ssh_pubkey
			EOF

        cat > "$STATE_DIR"/meta-data <<- EOF
			instance-id: node$i
			local-hostname: node$i
			EOF

        cat > "$STATE_DIR"/network-config <<- EOF
			network:
			  version: 2
			  ethernets:
			    eth0:
			      match:
			        macaddress: "52:54:00:00:00:$(printf '%02x' $i)"
			      dhcp4: false
			      addresses:
			        - 192.168.30.$((10 + i))/24
			EOF

        mkisofs -o "$STATE_DIR"/node$i-cloudinit.iso -V cidata -r -J "$STATE_DIR"/user-data "$STATE_DIR"/meta-data "$STATE_DIR"/network-config >/dev/null 2>&1

        rm "$STATE_DIR"/user-data
        rm "$STATE_DIR"/meta-data
        rm "$STATE_DIR"/network-config

        # Create main disk, backed by read-only cloud image
        qemu-img create -f qcow2 -F qcow2 -b "$BASE_IMAGE_FILENAME" "$STATE_DIR"/node$i.qcow2 20G >/dev/null 2>&1

        # Create SSH config to be used for accessing the VM
        cat >> "$SSH_CONFIG" <<- EOF
			Host node$i
			  StrictHostKeyChecking no
			  UserKnownHostsFile /dev/null
			  IdentityFile $STATE_DIR/id_ed25519
			  IdentitiesOnly yes
			  User root
			  Hostname localhost
			  Port $((3100 + i))
			EOF
    done

    start
}

case $1 in
    start)
        shift
        start "$@"
        ;;
    stop)
        shift
        stop "$@"
        ;;
    destroy)
        shift
        destroy "$@"
        ;;
    create)
        shift
        create "$@"
        ;;
esac