#!/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"
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=(
    [$X64]="/usr/share/edk2/ovmf/OVMF_CODE.fd"
    [$A64]="/usr/share/edk2/aarch64/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() {
    if [[ "$1" != "$X64" && "$1" != "$A64" ]]
    then
        echo "Usage $SCRIPT $OPTION $X64|$A64"
        return 1
    fi
    return 0
}

vm_boot() {
    local PID_FILE="${PID_FILES[$1]}"
    if [[ -f "$PID_FILE" ]]
    then
        echo "VM $1 is running"
        return
    fi
    local CMD=(
        "${QEMU[$1]}"
        "-M" "${MACHINES[$1]}"
        "-cpu" "max"
        "-m" "$MEMORY"
        "-smp" "$CORES"
        "-device" "$DEVICE"
        "-netdev" "user,id=net0,hostfwd=tcp::${SSH_PORTS[$1]}-:22"
        "-bios" "${BIOS[$1]}"
        "-hda" "${HDA[$1]}"
        "-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
    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"
        else
            echo "Invalid Machine Specified"
        fi
        return
    fi
    for NAME in "${NAMES[@]}"
    do
        vm_boot $NAME "-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
    fi
    local SSH_OPTIONS=()
    vm_ssh_options $MACHINE SSH_OPTIONS
    "${SSH_OPTIONS[@]}" "$USER@$HOST" $2
}

vm_ssh() {
    vm_dir_check
    ssh-add "$USER/$USER" &> /dev/null
    vm_cmd $VM
}

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

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

vm_test_cmd() {
    local MACHINE=$1
    vm_cmd $MACHINE "make kpl_test; ./kpl_test -a"
}

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 NAME in "${NAMES[@]}"
    do
        vm_rsync_machine $NAME
        vm_test_cmd $NAME
        TEST_STATUS[$NAME]=$?
    done
    for NAME in "${NAMES[@]}"
    do
        if [[ ${TEST_STATUS[$NAME]} != 0 ]]
        then
            echo "----------------------------"
            echo "  MACHINE $NAME FAILED TESTS"
            echo "----------------------------"
        fi
    done
}

vm_stop() {
    local ACTION="$1"
    vm_dir_check
    if [[ "$ACTION" != "KILL" ]]
    then
        ssh-add "$USER/$USER" &> /dev/null
    fi
    for NAME in "${NAMES[@]}"
    do
        local PID_FILE="${PID_FILES[$NAME]}"
        if [[ -f "$PID_FILE" ]]
        then
            if [[ "$ACTION" == "KILL" ]]
            then
                kill $(cat $PID_FILE)
                echo "VM $NAME force stop"
            else
                vm_cmd $NAME "sudo poweroff"
                echo "VM $NAME powering off"
            fi
        fi
   done
}

vm_clean() {
    if [[ ! -d "$DIR" ]]
    then
        return
    fi
    local CHOICE
    read -n 1 -r -s -p "Are you sure? y/n " CHOICE
    if [[ "$CHOICE" != "y" ]]
    then
        echo
        return
    fi
    vm_stop "KILL"
    rm -r $DIR
    echo -e "\nRemoved $DIR directory"
}

vm_help() {
    echo "See TESTING"
    echo "Usage $SCRIPT -init|-setup|-cli|-start|-ssh|-sync|-test|-stop|-kill|-clean|-help"
}

case $OPTION in
    -init)
        vm_init
        ;;
    -setup)
        vm_setup
        ;;
    -cli)
        vm_cli
        ;;
    -start)
        vm_start
        ;;
    -ssh)
        vm_ssh
        ;;
    -sync)
        vm_sync
        ;;
    -test)
        vm_test
        ;;
    -stop)
        vm_stop
        ;;
    -kill)
        vm_stop "KILL"
        ;;
    -clean)
        vm_clean
        ;;
    *)
        vm_help
        ;;
esac