#!/usr/bin/bash
#
# Copyright (C) 2024 Eugene 'Vindex' Stulin
# Distributed under the Boost Software License 1.0.

APP_NAME=vmm
APP_VERSION=0.1.0
[[ -z $LANG ]] && APP_LANG=en_US || APP_LANG=${LANG%.*}

set -eu -o pipefail

DATA_DIR="$(dirname -- "${BASH_SOURCE[0]}")/../share"
HELP_DIR=${DATA_DIR}/help
HELP_FILE=${HELP_DIR}/${APP_LANG}/${APP_NAME}/help.txt

SSH_CMD="ssh -o StrictHostKeyChecking=no"
SCP_CMD="scp -o StrictHostKeyChecking=no"

if [ ! -v "${VMM_CACHE_DIR+set}" ]; then
    VMM_CACHE_DIR=~/.cache/vmm
fi
JSON=vmmfile.json

readonly SIGS="EXIT HUP INT TERM ERR"

# The function prints help information.
Print_Help() {
    cat "${HELP_FILE}"
}


# The function prints the script version.
Print_Version() {
    echo "$APP_VERSION"
}


Wrong_Usage() {
    echo "Wrong usage. See: vmm --help" >&2
    exit 1
}


Check_JSON_Error() {
    if [[ "$1" == "null" || "$1" == "" ]]; then
        echo "JSON processing error. Missing field: \"$2\"." >&2
        exit 1
    fi
}


Read_JSON() {
    IM_SRC=$(jq -r -e '.base' "$JSON")
    Check_JSON_Error "$IM_SRC" "base"

    VM_DISK=$(jq -r '.vm_disk' "$JSON")
    Check_JSON_Error "$VM_DISK" "vm_disk"

    VM_NAME=$(jq -r '.name' "$JSON")
    Check_JSON_Error "$VM_NAME" "name"

    VM_MEM=$(jq -r '.memory' "$JSON")
    Check_JSON_Error "$VM_MEM" "memory"

    VM_CPUS=$(jq -r '.cpus' "$JSON")
    Check_JSON_Error "$VM_CPUS" "cpus"

    VM_OS=$(jq -r '.os' "$JSON")
    Check_JSON_Error "$VM_OS" "os"
}


readonly ATTEMPT_LIMIT=20

Wait_Access_And_Get_IP() {
    echo "Wait..."
    FOUND_IP=''
    declare -i COUNTER=0
    while [[ "$FOUND_IP" == "" && $((COUNTER < ATTEMPT_LIMIT)) ]]; do
        FOUND_LINE=$(virsh domifaddr "$VM_NAME" | head -3 | tail -1)
        FOUND_IP=$(echo "$FOUND_LINE" | awk '{ print $4 }')
        FOUND_IP=${FOUND_IP%/*}
        COUNTER=$((COUNTER+1))
        sleep 1
    done
    if [[ "$FOUND_IP" == "" ]]; then
        echo "IP not found for $VM_NAME" >&2
        exit 1
    fi
    mkdir -p "$HOME/.ssh/"
    local KNOWN_HOSTS="$HOME/.ssh/known_hosts"
    [[ ! -e "$KNOWN_HOSTS" ]] && touch "$KNOWN_HOSTS"
    ssh-keygen -f "$KNOWN_HOSTS" -R "$FOUND_IP"
    $SSH_CMD "root@$FOUND_IP" 'echo "VM is accessed."'
}


Destroy_And_Undefine() {
    virsh destroy "$VM_NAME"
    virsh undefine "$VM_NAME"
}


Initialize_VM() {
    local STOP_MODE="$1"

    mkdir -p "$VMM_CACHE_DIR"
    local IMAGE="$IM_SRC"
    if [[ "$IM_SRC" == *://* ]]; then  # URL
        local BASENAME="${IM_SRC##*/}"
        IMAGE="$VMM_CACHE_DIR/$BASENAME"
        if [[ ! -e "$IMAGE" ]]; then
            wget -P "$VMM_CACHE_DIR" "$IM_SRC"
        fi
    fi
    mkdir -p "$(dirname "$VM_DISK")"
    if [[ -e "$VM_DISK" ]]; then
        rm -f "$VM_DISK"
    fi
    cp "$IMAGE" "$VM_DISK"

    virt-install --name "$VM_NAME" \
        --memory "$VM_MEM" --vcpus="$VM_CPUS" --cpu=host \
        --disk "$VM_DISK",bus=virtio,cache=none,io=native --import \
        --noautoconsole --virt-type=kvm --osinfo "$VM_OS" \
        --network network=default
    trap Destroy_And_Undefine $SIGS
    if [[ $(jq 'has("init_commands")' "$JSON") == "false" ]]; then
        echo 'JSON processing error. Missing field: "init_commands".' >&2
        exit 1
    fi
    Wait_Access_And_Get_IP
    $SSH_CMD "root@$FOUND_IP" bash <<< $(jq -r -e '.init_commands[]' "$JSON")
    trap - $SIGS
    if [[ "$STOP_MODE" == TRUE ]]; then
        virsh destroy "$VM_NAME"
    fi
}
Parse_Args_And_Init() {
    shift
    local STOP_MODE=FALSE
    while getopts ':-:' VAL; do
        case $VAL in
            -)  # long arguments (with --)
                case $OPTARG in
                    stop) STOP_MODE=TRUE  ;;
                    *) echo "Unknown argument: --$OPTARG" ; exit 1
                esac
                ;;
            *) echo "Unknown argument: $OPTARG" ; exit 1 ;;
        esac
    done
    shift $((OPTIND-1))
    [[ $# -ne 0 ]] && Wrong_Usage
    Initialize_VM "$STOP_MODE"
}



Run_VM() {
    virsh start "$VM_NAME" || true
}


Stop_VM() {
    virsh destroy "$VM_NAME"
}


Eliminate_VM() {
    local RM_DISK_MODE="$1"
    Stop_VM 2>/dev/null || true
    virsh undefine "$VM_NAME"
    if [[ "$RM_DISK_MODE" == TRUE ]]; then
        rm -f "$VM_DISK"
    fi
}
Parse_Args_And_Eliminate() {
    shift
    local RM_DISK_MODE=FALSE
    while getopts ':-:' VAL; do
        case $VAL in
            -)  # long arguments (with --)
                case $OPTARG in
                    rm-disk) RM_DISK_MODE=TRUE  ;;
                    *) echo "Unknown argument: --$OPTARG" ; exit 1
                esac
                ;;
            *) echo "Unknown argument: $OPTARG" ; exit 1 ;;
        esac
    done
    shift $((OPTIND-1))
    [[ $# -ne 0 ]] && Wrong_Usage
    Eliminate_VM "$RM_DISK_MODE"
}


Run_Commands() {
    set +e
    if [[ $(jq 'has("commands")' "$JSON") == "false" ]]; then
        echo 'JSON processing error. Missing field: "commands".' >&2
        exit 1
    fi
    if [[ ! $(jq -e ".commands.\"$1\"[]" "$JSON" 2>/dev/null) ]]; then
        echo JSON processing error. Missing field: "commands[$1]". >&2
        exit 1
    fi
    $SSH_CMD "root@$FOUND_IP" bash <<< $(jq -r ".commands.\"$1\"[]" "$JSON")
    set -e
}


Run_Shell_Command() {
    $SSH_CMD "root@$FOUND_IP" "$1"
}


Parse_Args_And_Run_Command() {
    shift
    local RM_MODE=FALSE
    local CMD_NAME=""
    local SHELL_CMD=""
    local INIT_MODE=FALSE
    local RUN_MODE=FALSE
    local UPL_MODE=FALSE
    while getopts ':-:c:s:iru' VAL; do
        case $VAL in
            i) INIT_MODE=TRUE ;;
            r) RUN_MODE=TRUE ;;
            c) CMD_NAME="$OPTARG" ;;
            s) SHELL_CMD="$OPTARG" ;;
            u) UPL_MODE=TRUE ;;
            :)
                echo "No argument supplied ($VAL)" >&2
                exit 1
                ;;
            -)  # long arguments (with --)
                case $OPTARG in
                    init) INIT_MODE=TRUE ;;
                    run) RUN_MODE=TRUE ;;
                    rm) RM_MODE=TRUE  ;;
                    command=*) CMD_NAME="${OPTARG#*=}" ;;
                    command)
                        CMD_NAME="${!OPTIND}"
                        ((OPTIND++))
                        ;;
                    *) echo "Unknown argument: --$OPTARG" ; exit 1
                esac
                ;;
            *) echo "Unknown argument: $OPTARG" ; exit 1 ;;
        esac
    done

    shift $((OPTIND-1))
    [[ $# -ne 0 ]] && Wrong_Usage

    if [[ "$CMD_NAME" == "" && "$SHELL_CMD" == "" ]]; then
        echo "Command to execute is not defined." >&2
        exit 1
    fi
    if [[ "$INIT_MODE" == TRUE ]]; then
        Initialize_VM FALSE
    elif [[ "$RUN_MODE" == TRUE ]]; then
        Run_VM
    fi
    Wait_Access_And_Get_IP
    if [[ "$UPL_MODE" == TRUE ]]; then
        Upload . /vmm
    fi
    if [[ "$CMD_NAME" != "" ]]; then
        Run_Commands "$CMD_NAME"
    else
        Run_Shell_Command "$SHELL_CMD"
    fi
    if [[ "$RM_MODE" == TRUE ]]; then
        Eliminate_VM TRUE
    fi
}


Upload() {
    local SRC="$1"
    if [[ "$SRC" == "." ]]; then
        SRC="../${PWD##*/}"
    fi
    local DST="$2"
    ${SCP_CMD} -r "${SRC}" "root@$FOUND_IP:$DST" >/dev/null
}


Take() {
    local SRC="$1"
    local DST="$2"
    ${SCP_CMD} -r "root@$FOUND_IP:$SRC" "${DST}" >/dev/null
}


Parse_Args_And_Upload_or_Take() {
    local UP_MODE=TRUE
    if [[ "$1" == "take" ]]; then
        UP_MODE=FALSE
    fi
    shift
    local FROM=""
    local TO=""
    while getopts ':-:' VAL; do
        case $VAL in
            -)  # long arguments (with --)
                case $OPTARG in
                    from=*)
                        FROM="${OPTARG#*=}"
                        ;;
                    from)
                        FROM="${!OPTIND}"
                        ((OPTIND++))
                        ;;
                    to=*)
                        TO="${OPTARG#*=}"
                        ;;
                    to)
                        TO="${!OPTIND}"
                        ((OPTIND++))
                        ;;
                    *) echo "Unknown argument: --$OPTARG" ; exit 1
                esac
                ;;
            *) echo "Unknown argument: $OPTARG" ; exit 1 ;;
        esac
    done

    shift $((OPTIND-1))
    [[ $# -ne 0 ]] && Wrong_Usage

    if [[ "$FROM" == "" || "$TO" == "" ]]; then
        Wrong_Usage
    fi
    Wait_Access_And_Get_IP
    if [[ "$UP_MODE" == TRUE ]]; then
        Upload "$FROM" "$TO"
    else
        Take "$FROM" "$TO"
    fi
}



Main_Fn() {
    if [[ "$1" == "-h" || "$1" == "--help" ]]; then
        [[ $# -ne 1 ]] && Wrong_Usage
        Print_Help
        exit 0
    elif [[ "$1" == "-v" || "$1" == "--version" ]]; then
        [[ $# -ne 1 ]] && Wrong_Usage
        Print_Version
        exit 0
    elif [[ "$1" == "init" ]]; then
        Read_JSON
        Parse_Args_And_Init "$@"
    elif [[ "$1" == "run" ]]; then
        Read_JSON
        Run_VM
    elif [[ "$1" == "stop" ]]; then
        Read_JSON
        Stop_VM
    elif [[ "$1" == "eliminate" ]]; then
        Read_JSON
        Parse_Args_And_Eliminate "$@"
    elif [[ "$1" == "cmd" ]]; then
        Read_JSON
        Parse_Args_And_Run_Command "$@"
    elif [[ "$1" == "upload" || "$1" == "take" ]]; then
        Read_JSON
        Parse_Args_And_Upload_or_Take "$@"
    elif [[ "$1" == "status" ]]; then
        [[ $# -ne 1 ]] && Wrong_Usage
        Read_JSON
        virsh dominfo "$VM_NAME"
        virsh domifaddr "$VM_NAME"
    elif [[ "$1" == "clean-cache" ]]; then
        [[ $# -ne 1 ]] && Wrong_Usage
        rm -f "$VMM_CACHE_DIR"/*
    fi
}
Main_Fn "$@"
