#!/bin/bash

LIB_DIR="$( realpath "$(dirname "$(realpath "$0")")/../../eehyp-tools" )"
if [[ -e "$LIB_DIR/common" ]]; then
    # shellcheck source=/dev/null
    source "$LIB_DIR/common"
else
    echo "Failed to load eehyp-tools common lib"
    exit 1
fi

function usage () {
    ERROR="$1"
    EXIT_CODE=0
    if [[ -n "$ERROR" ]]; then
        echo "$ERROR"
        echo
        EXIT_CODE=3
    fi
    cat << EOF
Usage : $0 [-d] [-h] [options]
    -r hyp_reserved_mem    Specify hypervisor reserved memory (in Mo, default: ${HYP_RESERVED_MEMORY})
    -w warning_cpu_ratio   Warning threshold ratio of allocated vCPU (default: ${WARNING_CPU_THRESHOLD})
    -c critical_cpu_ratio  Critical threshold ratio of allocated vCPU (default: ${CRITICAL_CPU_THRESHOLD})
    -d                     Debug mode
    -h                     Show this message
EOF
    exit $EXIT_CODE
}

DEBUG=0
while getopts "hr:w:c:d" OPTION; do
    case $OPTION in
        r)
            HYP_RESERVED_MEMORY=$OPTARG
        ;;
        w)
            WARNING_CPU_THRESHOLD=$OPTARG
        ;;
        c)
            CRITICAL_CPU_THRESHOLD=$OPTARG
        ;;
        d)
            # shellcheck disable=SC2034
            DEBUG=1
        ;;
        h)
            usage
        ;;
        \?)
            usage "Invalid option $OPTION"
    esac
done

HYP_VM=0
HYP_ALLOCATED_CPU=0
HYP_ALLOCATED_MEMORY=0
set_newline_ifs
for line in $( pgrep -f '(qemu-system|/usr/bin/kvm)' -a ); do
    # shellcheck disable=SC2001
    NAME=$( sed 's/^.* -name \([^ ]*\) .*$/\1/' <<< "$line" | sed 's/guest=\([^,]*\),.*/\1/' )
    (( HYP_VM+=1 ))

    # Retrieve allocated memories
    # shellcheck disable=SC2001
    ALLOCATED_MEMORY=$( sed 's/^.* -m \([0-9]*\) .*$/\1/' <<< "$line" )
    (( HYP_ALLOCATED_MEMORY+=ALLOCATED_MEMORY ))

    # Retrieve allocated CPU
    # shellcheck disable=SC2001
    VCPU=$( sed 's/^.* -smp \([0-9]*\),.*$/\1/' <<< "$line" )
    (( HYP_ALLOCATED_CPU+=VCPU ))
    debug "VM $NAME => Mem=$ALLOCATED_MEMORY / vCPU=$VCPU"
done
restore_ifs

debug "Total hypervisor allocated memory : ${HYP_ALLOCATED_MEMORY}Mo"
debug "Total hypervisor allocated vCPU : ${HYP_ALLOCATED_CPU}"

HYP_MEMORY=$( grep '^MemTotal:' < /proc/meminfo | awk '{print $2}' )
(( HYP_MEMORY=HYP_MEMORY/1024 ))
debug "Hypervisor total memory: $HYP_MEMORY Mo"

(( HYP_ALLOCABLE_MEMORY=HYP_MEMORY-HYP_RESERVED_MEMORY ))
debug "Hypervisor allocable memory: ${HYP_ALLOCABLE_MEMORY}Mo"

(( HYP_ALLOCATED_MEMORY_PERC=HYP_ALLOCATED_MEMORY*100/HYP_ALLOCABLE_MEMORY ))
debug "Hypervisor allocated memory ratio: $HYP_ALLOCATED_MEMORY_PERC%"

HYP_CPU=$( grep -c ^processor < /proc/cpuinfo )
debug "Hypervisor total CPU thread: $HYP_CPU"

(( HYP_ALLOCATED_CPU_PERC=HYP_ALLOCATED_CPU*100/HYP_CPU ))
debug "Hypervisor allocated CPU ratio: $HYP_ALLOCATED_CPU_PERC%"

EXIT_CODE=0
EXIT_STATUS=OK
if [[ $HYP_ALLOCATED_MEMORY -ge $HYP_MEMORY ]]; then
    EXIT_CODE=2
    EXIT_STATUS=CRITICAL
elif [[ $HYP_ALLOCATED_MEMORY -ge $HYP_ALLOCABLE_MEMORY ]]; then
    EXIT_CODE=1
    EXIT_STATUS=WARNING
fi

(( HYP_WARNING_ALLOCATED_CPU_THRESHOLD=HYP_CPU*WARNING_CPU_THRESHOLD ))
(( HYP_CRITICAL_ALLOCATED_CPU_THRESHOLD=HYP_CPU*CRITICAL_CPU_THRESHOLD ))
debug "Hypervisor critical/warning allocated CPU thresholds:" \
    "$HYP_WARNING_ALLOCATED_CPU_THRESHOLD / $HYP_CRITICAL_ALLOCATED_CPU_THRESHOLD"

if [[ $HYP_ALLOCATED_CPU -ge $HYP_CRITICAL_ALLOCATED_CPU_THRESHOLD ]]; then
    EXIT_CODE=2
        EXIT_STATUS=CRITICAL
elif [[ $EXIT_CODE -le 1 ]] && \
    [[ $HYP_ALLOCATED_CPU -ge $HYP_WARNING_ALLOCATED_CPU_THRESHOLD ]]; then
        EXIT_CODE=1
        EXIT_STATUS=WARNING
fi

add_perfdata \
    allocated_memory \
    ${HYP_ALLOCATED_MEMORY}Mo \
    $HYP_ALLOCABLE_MEMORY $HYP_MEMORY \
    0 $HYP_MEMORY

add_perfdata \
    allocated_vcpu \
    ${HYP_ALLOCATED_CPU} \
    $HYP_WARNING_ALLOCATED_CPU_THRESHOLD $HYP_CRITICAL_ALLOCATED_CPU_THRESHOLD \
    0 "$HYP_CPU"

echo "$EXIT_STATUS - $HYP_ALLOCATED_MEMORY_PERC% allocated memory" \
    "($HYP_ALLOCATED_MEMORY/${HYP_ALLOCABLE_MEMORY}Mo allocable, ${HYP_MEMORY}Mo total)," \
    "${HYP_ALLOCATED_CPU} allocated vCPU (${HYP_WARNING_ALLOCATED_CPU_THRESHOLD} warning /" \
    "${HYP_CRITICAL_ALLOCATED_CPU_THRESHOLD} critical) | $(format_perfdata)"

exit $EXIT_CODE
