#!/bin/bash

# Copyright (C) 2024-2025 Eugene 'Vindex' Stulin
# Distributed under the Boost Software License 1.0.

APP_NAME=qmm
APP_VERSION=0.2.6
[[ -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
if [[ ! -e "$HELP_FILE" ]]; then
    HELP_FILE="$(dirname -- "${BASH_SOURCE[0]}")/../help/${APP_LANG}.txt"
fi

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

if [ ! -v "${QMM_CACHE_DIR+set}" ]; then
    QMM_CACHE_DIR=~/.cache/qmm
fi
JSON=qmmfile.json

readonly SIGS="EXIT HUP INT TERM ERR"

STORED_ERR_CODE=0


# 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: qmm --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_Field() {
    local FIELD="$1"
    local VALUE=$(jq -r -e ".${FIELD}" "$JSON")
    Check_JSON_Error "${VALUE}" "${FIELD}"
    echo "${VALUE}"
}


Read_JSON() {
    IM_SRC="$(Read_JSON_Field 'base')"
    VM_DISK="$(Read_JSON_Field 'vm_disk')"
    VM_NAME="$(Read_JSON_Field 'name')"
    VM_MEM="$(Read_JSON_Field 'memory')"
    VM_CPUS="$(Read_JSON_Field 'cpus')"
    VM_OS="$(Read_JSON_Field 'os')"

    SH=$(jq -r '.interpreter' "$JSON")
    if [[ "$SH" == "" || "$SH" == "null" ]]; then
        SH="/bin/bash"
    fi
}


Label_With_Border() {
    local LABEL=$1
    local EDGE='--------------------------------------------------'
    echo "$EDGE"
    echo ">: $LABEL"
    echo "$EDGE"
}


readonly ATTEMPT_LIMIT=60

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 "$QMM_CACHE_DIR"
    local IMAGE="$IM_SRC"
    if [[ "$IM_SRC" == *://* ]]; then  # URL
        local BASENAME="${IM_SRC##*/}"
        IMAGE="$QMM_CACHE_DIR/$BASENAME"
        if [[ ! -e "$IMAGE" ]]; then
            wget -P "$QMM_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-model \
        --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
    if [[ $(jq -r -e '.init_commands[]' "$JSON") == "" ]]; then
        echo "There no initialization commands."
        trap - $SIGS
        return
    fi
    Wait_Access_And_Get_IP
    echo ""
    Label_With_Border "VM initialization"
    $SSH_CMD "root@$FOUND_IP" "$SH" <<< $(jq -r -e '.init_commands[]' "$JSON")
    trap - $SIGS
    if [[ "$STOP_MODE" == TRUE ]]; then
        virsh destroy "$VM_NAME"
    fi
    Label_With_Border "End of initialization" && echo ""
}
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"
}


Connect_With_SSH() {
    if virsh dominfo "$VM_NAME" ; then
        Wait_Access_And_Get_IP
        $SSH_CMD "root@$FOUND_IP"
    else
        echo "VM is not running"
    fi
}


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


Stop_VM() {
    local output=$(virsh destroy "$VM_NAME")
    if [[ "$output" != "" ]]; then
        echo "$output"
    fi
}


Eliminate_VM() {
    local RM_DISK_MODE="$1"
    Stop_VM 2>/dev/null || true
    local output=$(virsh undefine "$VM_NAME" 2>&1)
    echo "$output"
    if [[ "$RM_DISK_MODE" == TRUE ]]; then
        rm -f "$VM_DISK"
        echo "The virtual disk \"$VM_DISK\" has been removed"
    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
    local CMD=$(jq -r ".commands.\"$1\"[]" "$JSON")
    $SSH_CMD "root@$FOUND_IP" "$CMD"
    STORED_ERR_CODE=$?
    set -e
}


Run_Shell_Command() {
    set +e
    $SSH_CMD "root@$FOUND_IP" "$1"
    STORED_ERR_CODE=$?
    set -e
}


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

    if [[ $(jq 'has("commands")' "$JSON") == "false" ]]; then
        echo 'JSON processing error. Missing field: "commands".' >&2
        exit 1
    fi
    if [[ "$CMD_NAME" == "" ]]; then
        :
    elif [[ ! $(jq -e ".commands.\"$CMD_NAME\"[]" "$JSON" 2>/dev/null) ]]; then
        echo JSON processing error. Missing field: "commands['$CMD_NAME']". >&2
        echo "The '$CMD_NAME' command does not exist."
        exit 1
    fi

    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

    # !!! catch signals
    [[ "$INIT_MODE" == TRUE ]] && trap Destroy_And_Undefine $SIGS
    # !!!

    Wait_Access_And_Get_IP
    if [[ "$UPL_MODE" == TRUE ]]; then
        Upload . /qmm
    fi

    echo "" && Label_With_Border "VM output"
    if [[ "$CMD_NAME" != "" ]]; then
        Run_Commands "$CMD_NAME"
    else
        Run_Shell_Command "${SHELL_CMD}"
    fi
    Label_With_Border "End of output" && echo ""

    if [[ "${STORED_ERR_CODE}" != "0" ]]; then
        echo -e "The requested command completed in failure.\n"
        Eliminate_VM TRUE
    fi

    # !!! end of trap
    [[ "$INIT_MODE" == TRUE ]] && trap - $SIGS
    # !!!

    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
}


Show_Status() {
    local OUTPUT=$(virsh dominfo "$VM_NAME" 2>&1)
    echo "$OUTPUT"
    if [[ "$OUTPUT" != error* ]]; then
        local OUTPUT2=$(virsh domifaddr "$VM_NAME" 2>&1)
        echo "$OUTPUT2"
        if [[ "$OUTPUT2" == error* ]]; then
            exit 1
        fi
    else
        exit 1
    fi
}


Main_Fn() {
    if [[ $# -eq 0 ]]; then
        Wrong_Usage
    elif [[ "$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" == "connect" ]]; then
        Read_JSON
        Connect_With_SSH
    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
        Show_Status
    elif [[ "$1" == "clean-cache" ]]; then
        [[ $# -ne 1 ]] && Wrong_Usage
        rm -f "$QMM_CACHE_DIR"/*
    else
        Wrong_Usage
    fi
}
Main_Fn "$@"
