#!/bin/bash

SCRIPT="$0"
OPTION="$1"
VM="$2"
ARG_COUNT=$#
DIR="qemu"
USER="user"
HOST="localhost"
RSYNC_LIST="rsync_list"
CORES="2"
MEMORY="2G"
DEVICE="virtio-net-pci,netdev=net0"
HDA_URL="https://cloud.debian.org/images/cloud/trixie/latest"
TEST_COMMAND="make kpl_test && arch && ./kpl_test -a"
CLEAN_COMMAND="make clean; find -type f -regex '^.*\.[ch]' -exec rm -v {} \;"
X64="X64"
A64="A64"
NAMES=("$X64" "$A64")
declare -A QEMU=(
    [$X64]="qemu-system-x86_64"
    [$A64]="qemu-system-aarch64"
)
declare -A HDA=(
    [$X64]="debian-13-nocloud-amd64.qcow2"
    [$A64]="debian-13-nocloud-arm64.qcow2"
)
declare -A BIOS_FD=(
    [$X64]="OVMF_CODE.fd"
    [$A64]="QEMU_EFI.fd"
)
declare -A SSH_PORTS=(
    [$X64]="2222"
    [$A64]="2223"
)
declare -A PID_FILES=(
    [$X64]="$X64.pid"
    [$A64]="$A64.pid"
)
declare -A MACHINES=(
    [$X64]="pc"
    [$A64]="virt"
)

vm_init() {
   if [[ ! -d "$DIR" ]]
   then
        mkdir $DIR
        echo "DO NOT MODIFY CONTENTS OF THIS DIRECTORY" > "$DIR/README"
   fi
   make rsync_list > "$DIR/$RSYNC_LIST"
   cd $DIR
   if [[ ! -d "$USER" ]]
   then
        mkdir $USER
        ssh-keygen -t ed25519 -f "$USER/$USER" -N ""
   fi
   for H in ${HDA[@]}
   do
        if [[ ! -f "$H" ]]
        then
            wget "$HDA_URL/$H"
        fi
   done
}

vm_setup() {
    local PUB=$(cat "$DIR/$USER/$USER.pub")
    local HEADER="COPY AND RUN IN QEMU:"
    echo $HEADER
    echo "    sed -i 's/GRUB_TIMEOUT=5/GRUB_TIMEOUT_STYLE=hidden\nGRUB_TIMEOUT=0/' /etc/default/grub; update-grub; sed -i 's/^%sudo\s\+ALL=(ALL:ALL)\sALL$/%sudo   ALL=(ALL:ALL) NOPASSWD:ALL/' /etc/sudoers; adduser --disabled-password --gecos "" --shell /bin/bash $USER; adduser $USER sudo; su - $USER"
    echo $HEADER
    echo "    mkdir .ssh; cd .ssh; echo \"$PUB\" > authorized_keys; chmod 600 authorized_keys; cd ~"
    echo $HEADER
    echo "    sudo apt update; yes | sudo apt upgrade; yes | sudo apt install ssh make gcc gdb rsync; sudo poweroff"
}

vm_machine_check() {
    local MACHINE=$1
    if [[ "$MACHINE" != "$X64" && "$MACHINE" != "$A64" ]]
    then
        echo "Usage $SCRIPT $OPTION $X64|$A64"
        return 1
    fi
    return 0
}

vm_boot() {
    local MACHINE=$1
    local PID_FILE="${PID_FILES[$MACHINE]}"
    if [[ -f "$PID_FILE" ]]
    then
        echo "VM $MACHINE is running"
        return 1
    fi
    local BIOS=$(locate "${BIOS_FD[$MACHINE]}" | grep "edk2" | head -n 1)
    local CMD=(
        "${QEMU[$MACHINE]}"
        "-M" "${MACHINES[$MACHINE]}"
        "-cpu" "max"
        "-m" "$MEMORY"
        "-smp" "$CORES"
        "-device" "$DEVICE"
        "-netdev" "user,id=net0,hostfwd=tcp::${SSH_PORTS[$MACHINE]}-:22"
        "-bios" "$BIOS"
        "-hda" "${HDA[$MACHINE]}"
        "-pidfile" "$PID_FILE"
    )
    for ARG in "${@:2}"
    do
        CMD+=("$ARG")
    done
    echo "${CMD[*]}"
    "${CMD[@]}"
}

vm_dir_check() {
    if [[ ! -d "$DIR" ]]
    then
        echo "Environment not setup, run: $SCRIPT -init"
        exit 1
    fi
    cd $DIR
}

vm_cli() {
    vm_dir_check
    vm_machine_check $VM
    if [[ $? != 0 ]]
    then
        return 1
    fi
    vm_boot $VM "-nographic"
}

vm_start() {
    vm_dir_check
    if [[ $ARG_COUNT == 2 ]]
    then
        vm_machine_check $VM
        if [[ $? == 0 ]]
        then
            vm_boot $VM "-display" "none" "-daemonize"
            return
        else
            echo "Invalid Machine Specified"
            return 1
        fi
    fi
    for MACHINE in "${NAMES[@]}"
    do
        vm_boot $MACHINE "-display" "none" "-daemonize"
    done
}

vm_ssh_options() {
    local MACHINE=$1
    local -n OPTIONS=$2
    OPTIONS=(
        "ssh"
        "-o" "UserKnownHostsFile=/dev/null"
        "-o" "StrictHostKeyChecking=no"
        "-p" "${SSH_PORTS[$MACHINE]}"
    )
}

vm_cmd() {
    local MACHINE=$1
    vm_machine_check $MACHINE
    if [[ $? != 0 ]]
    then
        return 1
    fi
    local SSH_OPTIONS=()
    vm_ssh_options $MACHINE SSH_OPTIONS
    "${SSH_OPTIONS[@]}" "$USER@$HOST" $2
}

vm_ssh_add() {
    ssh-add "$USER/$USER" &> /dev/null
}

vm_ssh() {
    vm_dir_check
    vm_ssh_add
    vm_cmd $VM
}

vm_rsync_machine() {
    local MACHINE=$1
    vm_machine_check $MACHINE
    if [[ $? != 0 ]]
    then
        echo "Invalid Machine Specified"
        return 1
    fi
    local SSH_OPTIONS=()
    vm_ssh_options $MACHINE SSH_OPTIONS
    local CMD=(
        "rsync"
        "-avz"
        "-e" "${SSH_OPTIONS[*]}"
        "--include-from=$DIR/$RSYNC_LIST"
        "--exclude=*"
        "$PWD/"
        "$USER@$HOST:/home/$USER"
    )
    vm_ssh_add
    "${CMD[@]}"
}

vm_sync() {
    if [[ "$ARG_COUNT" == 2 ]]
    then
        vm_rsync_machine $VM
        return
    fi
    for MACHINE in "${NAMES[@]}"
    do
        vm_rsync_machine $MACHINE
    done
}

vm_test_cmd() {
    local MACHINE=$1
    vm_cmd $MACHINE $TEST_COMMAND
}

vm_test() {
    if [[ "$ARG_COUNT" == 2 ]]
    then
        vm_rsync_machine $VM
        vm_test_cmd $VM
        return
    fi
    local -A TEST_STATUS=([$X64]=0 [$A64]=0)
    for MACHINE in "${NAMES[@]}"
    do
        vm_rsync_machine $MACHINE
        vm_test_cmd $MACHINE
        TEST_STATUS[$MACHINE]=$?
    done
    for MACHINE in "${NAMES[@]}"
    do
        if [[ "${TEST_STATUS[$MACHINE]}" != 0 ]]
        then
            echo "-----------------------"
            echo "  VM $MACHINE FAILED TESTS"
            echo "-----------------------"
        fi
    done
    return (( ${TEST_STATUS[$X64]} | ${TEST_STATUS[$A64]} ))
}

vm_stop() {
    vm_dir_check
    vm_ssh_add
    for MACHINE in "${NAMES[@]}"
    do
        local PID_FILE="${PID_FILES[$MACHINE]}"
        if [[ -f "$PID_FILE" ]]
        then
            vm_cmd $MACHINE "sudo poweroff"
            echo "VM $MACHINE powering off"
        fi
   done
}

vm_kill() {
    vm_dir_check
    for MACHINE in "${NAMES[@]}"
    do
        local PID_FILE="${PID_FILES[$MACHINE]}"
        if [[ -f "$PID_FILE" ]]
        then
            kill $(cat $PID_FILE)
            echo "VM $MACHINE force stop"
        fi
   done
}

vm_clean_cmd() {
    local MACHINE=$1
    vm_machine_check $MACHINE
    if [[ $? != 0 ]]
    then
        echo "Invalid Machine Specified"
        return 1
    fi
    vm_cmd $MACHINE $CLEAN_COMMAND
}

vm_clean() {
    vm_ssh_add
    if [[ "$ARG_COUNT" == 2 ]]
    then
        vm_clean_cmd $VM
        return
    fi
    for MACHINE in "${NAMES[@]}"
    do
        vm_clean_cmd $MACHINE
    done
}

vm_destroy() {
    if [[ ! -d "$DIR" ]]
    then
        return
    fi
    local CHOICE
    read -n 1 -r -p "Are you sure? [y/n] " CHOICE
    if [[ "$CHOICE" != "y" ]]
    then
        echo
        return
    fi
    echo
    vm_kill
    cd ..
    rm -rf $DIR
    echo -e "Removed $DIR directory"
}

declare -A SCRIPT_COMMANDS=(
    ["-init"]="vm_init"
    ["-setup"]="vm_setup"
    ["-cli"]="vm_cli"
    ["-start"]="vm_start"
    ["-ssh"]="vm_ssh"
    ["-sync"]="vm_sync"
    ["-test"]="vm_test"
    ["-clean"]="vm_clean"
    ["-stop"]="vm_stop"
    ["-kill"]="vm_kill"
    ["-destroy"]="vm_destroy"
    ["-help"]="vm_help"
)

vm_help() {
    local LENGTH="${#SCRIPT_COMMANDS[@]}"
    local I=0
    echo "See TESTING"
    echo -n "Usage $SCRIPT ["
    for CMD in "${!SCRIPT_COMMANDS[@]}"
    do
        echo -n "$CMD"
        if (( I + 1 < LENGTH ))
        then
            echo -n "|"
        else
            echo "]"
        fi
        (( I++ ))
    done
    return 2
}

vm_main() {
    for CMD in "${!SCRIPT_COMMANDS[@]}"
    do
        if [[ "$OPTION" == "$CMD" ]]
        then
            "${SCRIPT_COMMANDS[$CMD]}"
            return
        fi
    done
    echo -n "Invalid [OPTION], "
    vm_help
}

vm_main