#! /bin/bash
# vim: shiftwidth=4 tabstop=4 expandtab

HYP_STATS_STATUS=( all active inactive persistent transient running paused shutoff other )

# Auto-complete twik: allow sourcing this file from complete command to get supported args
[[ ${IN_AUTOCOMPLETER:-0} -eq 1 ]] && return

function get_hyp_stats() {
    local hyp=$1
    meminfo=$( run_on_hyp "$hyp" cat /proc/meminfo ) && \
    domstats=$( run_on_hyp "$hyp" "$REAL_VIRSH" domstats --raw --state --balloon --vcpu "${@:2}" ) && \
    cpu=$( run_on_hyp "$hyp" grep -c ^processor /proc/cpuinfo ) && \
    jq --null-input -r \
        --arg meminfo "$meminfo" \
        --arg domstats "$domstats" \
        --argjson cpu "$cpu" \
        --argjson vcpu_reserved "$HYP_RESERVED_VCPUS" \
        --argjson HYP_RESERVED_MEMORY $HYP_RESERVED_MEMORY \
        --argjson VCPU_OVER_PROVISIONING $VCPU_OVER_PROVISIONING \
        --argjson status_mapping '{
            "1": "running",
            "2": "blocked",
            "3": "paused",
            "4": "shutdown",
            "5": "shut off",
            "6": "crashed",
            "7": "pmsuspended"
        }' \
        '
        # Parse "virsh domstats" result
        reduce ($domstats | split("\n"))[] as $line (
            { vms: {}, current_vm: null };
            if $line | startswith("Domain:") then
                .current_vm = ($line | capture("Domain: '"'"'(?<vm>.+)'"'"'") as $m | $m.vm) |
                .vms[.current_vm] = {}
            elif .current_vm != null and ($line | test("[a-zA-Z0-9_.]+=")) then
                . |
                ($line | capture("(?<key>[a-zA-Z0-9_.]+)=(?<value>.+)")) as $kv |
                setpath(
                    ["vms", .current_vm] + ($kv.key | split("."));
                    ($kv.value | tonumber? // $kv.value)
                )
            else
                .
            end
        ) |
        # Remove temporary .current_vm key use during parsing process
        del(.current_vm) |
        # Add some keys on each VMs
        .vms |= with_entries(
            .value.status = ($status_mapping[.value.state.state | tostring] // .value.state.state) |
            .value.balloon.unused //= 0 |
            .value.balloon.used = (.value.balloon.maximum - .value.balloon.unused) |
            .value.balloon.used_percent = (
                if .value.balloon.current == 0 then 0 else (
                    .value.balloon.used * 100 / .value.balloon.current * 10 | floor / 10
                ) end
            )
        ) |
        # Parse meminfo
        .meminfo = reduce (
            $meminfo | split("\n") | .[] | capture("(?<key>[^:]+):\\s*(?<value>\\d+)")
        ) as {$key, $value} (
            {}; .[$key] = ($value | tonumber)
        ) |
        # Compute global info
        .memory = .meminfo.MemTotal |
        .memory_cached = .meminfo.Cached |
        .memory_cached_percent = (
            if .memory == 0 then 0 else (.memory_cached * 100 / .memory * 10 | floor / 10) end
        ) |
        .memory_buffers = .meminfo.Buffers |
        .memory_free = .meminfo.MemFree |
        .memory_free_percent = (
            if .memory == 0 then 0 else (.memory_free * 100 / .memory * 10 | floor / 10) end
        ) |
        .memory_used = (.memory - .memory_cached - .memory_buffers - .memory_free) |
        .memory_used_percent = (
            if .memory == 0 then 0 else (.memory_used * 100 / .memory * 10 | floor / 10) end
        ) |
        .memory_reserved = ($HYP_RESERVED_MEMORY * 1024) |
        .memory_allocatable = (.memory - .memory_reserved) |
        .memory_allocated = 0 + ([.vms[] | .balloon.current] | add) |
        .memory_allocated_percent = (
            if .memory_allocatable == 0 then 0 else (
                .memory_allocated * 100 / .memory_allocatable * 10 | floor / 10
            ) end
        ) |
        .memory_allocated_used = 0 + ([.vms[] | .balloon.used] | add) |
        .memory_allocated_used_percent = (
            if .memory_allocated_used == 0 then 0 else (
                .memory_allocated_used * 100 / .memory_allocated_used * 10 | floor / 10
            ) end
        ) |
        .memory_max_allocated = 0 + ([.vms[] | .balloon.maximum] | add) |
        .memory_max_allocated_percent = (
            if .memory_allocatable == 0 then 0 else (
                .memory_max_allocated * 100 / .memory_allocatable * 10 | floor / 10
            ) end
        ) |
        .memory_available = (.memory_allocatable - .memory_allocated) |
        .memory_available_percent = (
            if .memory_allocatable == 0 then 0 else (
                .memory_available * 100 / .memory_allocatable * 10 | floor / 10
            ) end
        ) |
        .swap = .meminfo.SwapTotal |
        .swap_cached = .meminfo.SwapCached |
        .swap_cached_percent = (
            if .swap == 0 then 0 else (.swap_cached * 100 / .swap * 10 | floor / 10) end
        ) |
        .swap_free = .meminfo.SwapFree |
        .swap_free_percent = (
            if .swap == 0 then 0 else (.swap_free * 100 / .swap * 10 | floor / 10) end
        ) |
        .swap_used = (.swap - .swap_cached - .swap_free) |
        .swap_used_percent = (
            if .swap == 0 then 0 else (.swap_used * 100 / .swap * 10 | floor / 10) end
        ) |
        .total_memory = (.memory + .swap) |
        .total_memory_free = (.memory_free + .swap_free) |
        .total_memory_free_percent = (
            if .total_memory == 0 then 0 else (
                .total_memory_free * 100 / .total_memory * 10 | floor / 10
            ) end
        ) |
        .total_memory_cached = (.memory_cached + .swap_cached) |
        .total_memory_cached_percent = (
            if .total_memory == 0 then 0 else (
                .total_memory_cached * 100 / .total_memory * 10 | floor / 10
            ) end
        ) |
        .total_memory_used = (.memory_used + .swap_used) |
        .total_memory_used_percent = (
            if .total_memory == 0 then 0 else (
                .total_memory_used * 100 / .total_memory * 10 | floor / 10
            ) end
        ) |
        .total_memory_used_percent_real_ram = (
            if .memory == 0 then 0 else (.total_memory_used * 100 / .memory * 10 | floor / 10) end
        ) |
        .cpu = $cpu |
        .vcpu_reserved = $vcpu_reserved |
        .vcpu_allocatable = ((.cpu * $VCPU_OVER_PROVISIONING) - .vcpu_reserved) |
        .vcpu_allocated = 0 + ([.vms[] | .vcpu.current] | add) |
        .vcpu_allocated_percent = (
            if .vcpu_allocatable == 0 then 0 else (
                .vcpu_allocated * 100 / .vcpu_allocatable * 10 | floor / 10
            ) end
        ) |
        .vcpu_max_allocated = 0 + ([.vms[] | .vcpu.maximum] | add) |
        .vcpu_max_allocated_percent = (
            if .vcpu_allocatable == 0 then 0 else (
                .vcpu_max_allocated * 100 / .vcpu_allocatable * 10 | floor / 10
            ) end
        ) |
        .vcpu_available = (.vcpu_allocatable - .vcpu_allocated) |
        .vcpu_available_percent = (
            if .vcpu_allocatable == 0 then 0 else (
                .vcpu_available * 100 / .vcpu_allocatable * 10 | floor / 10
            ) end
        ) |
        # Remove temporary meminfo key
        del(.meminfo)'
}

function get_cluster_stats() {
    local stats=$( jo hyps={} vms={} )
    for hyp in $HYPS; do
        hyp_stats=$( get_hyp_stats "$hyp" "${@:1}" ) && \
        stats=$(
            jq -r --arg hyp "$hyp" --argjson stats "$hyp_stats" '
                .hyps[$hyp] = $stats |
                del(.hyps[$hyp].vms) |
                .hyps[$hyp].vm_count = ($stats.vms | length) |
                .vms = (.vms * ($stats.vms | with_entries(.value.hypervisor = $hyp)))

            ' <<< "$stats"
        )
    done

    jq -r '. |
        .memory = 0 + ([.hyps[] | .memory] | add) |
        .memory_cached = 0 + ([.hyps[] | .memory_cached] | add) |
        .memory_cached_percent = (
            if .memory == 0 then 0 else (
                .memory_cached * 100 / .memory * 10 | floor / 10
            ) end
        ) |
        .memory_buffers = 0 + ([.hyps[] | .memory_buffers] | add) |
        .memory_free = 0 + ([.hyps[] | .memory_free] | add) |
        .memory_free_percent = (
            if .memory == 0 then 0 else (.memory_free * 100 / .memory * 10 | floor / 10) end
        ) |
        .memory_used = 0 + ([.hyps[] | .memory_used] | add) |
        .memory_used_percent = (
            if .memory == 0 then 0 else (.memory_used * 100 / .memory * 10 | floor / 10) end
        ) |
        .memory_reserved = 0 + ([.hyps[] | .memory_reserved] | add) |
        .memory_allocatable = 0 + ([.hyps[] | .memory_allocatable] | add) |
        .memory_allocated = 0 + ([.hyps[] | .memory_allocated] | add) |
        .memory_allocated_percent = (
            if .memory_allocatable == 0 then 0 else (
                .memory_allocated * 100 / .memory_allocatable * 10 | floor / 10
            ) end
        ) |
        .memory_allocated_used = 0 + ([.hyps[] | .memory_allocated_used] | add) |
        .memory_allocated_used_percent = (
            if .memory_allocated_used == 0 then 0 else (
                .memory_allocated_used * 100 / .memory_allocated_used * 10 | floor / 10
            ) end
        ) |
        .memory_max_allocated = 0 + ([.hyps[] | .memory_max_allocated] | add) |
        .memory_max_allocated_percent = (
            if .memory_allocatable == 0 then 0 else (
                .memory_max_allocated * 100 / .memory_allocatable * 10 | floor / 10
            ) end
        ) |
        .memory_available = 0 + ([.hyps[] | .memory_available] | add) |
        .memory_available_percent = (
            if .memory_allocatable == 0 then 0 else (
                .memory_available * 100 / .memory_allocatable * 10 | floor / 10
            ) end
        ) |
        .swap = 0 + ([.hyps[] | .swap] | add) |
        .swap_cached = 0 + ([.hyps[] | .swap_cached] | add) |
        .swap_cached_percent = (
            if .swap == 0 then 0 else (.swap_cached * 100 / .swap * 10 | floor / 10) end
        ) |
        .swap_free = 0 + ([.hyps[] | .swap_free] | add) |
        .swap_free_percent = (
            if .swap == 0 then 0 else (.swap_free * 100 / .swap * 10 | floor / 10) end
        ) |
        .swap_used = 0 + ([.hyps[] | .swap_used] | add) |
        .swap_used_percent = (
            if .swap == 0 then 0 else (.swap_used * 100 / .swap * 10 | floor / 10) end
        ) |
        .total_memory = 0 + ([.hyps[] | .total_memory] | add) |
        .total_memory_free = 0 + ([.hyps[] | .total_memory_free] | add) |
        .total_memory_free_percent = (
            if .total_memory == 0 then 0 else (
                .total_memory_free * 100 / .total_memory * 10 | floor / 10
            ) end
        ) |
        .total_memory_cached = 0 + ([.hyps[] | .total_memory_cached] | add) |
        .total_memory_cached_percent = (
            if .total_memory == 0 then 0 else (
                .total_memory_cached * 100 / .total_memory * 10 | floor / 10
            ) end
        ) |
        .total_memory_used = 0 + ([.hyps[] | .total_memory_used] | add) |
        .total_memory_used_percent = (
            if .total_memory == 0 then 0 else (
                .total_memory_used * 100 / .total_memory * 10 | floor / 10
            ) end
        ) |
        .total_memory_used_percent_real_ram = (
            if .memory == 0 then 0 else (.total_memory_used * 100 / .memory * 10 | floor / 10) end
        ) |
        .cpu = 0 + ([.hyps[] | .cpu] | add) |
        .vcpu_reserved = 0 + ([.hyps[] | .vcpu_reserved] | add) |
        .vcpu_allocatable = 0 + ([.hyps[] | .vcpu_allocatable] | add) |
        .vcpu_allocated = 0 + ([.hyps[] | .vcpu_allocated] | add) |
        .vcpu_allocated_percent = (
            if .vcpu_allocatable == 0 then 0 else (
                .vcpu_allocated * 100 / .vcpu_allocatable * 10 | floor / 10
            ) end
        ) |
        .vcpu_max_allocated = 0 + ([.hyps[] | .vcpu_max_allocated] | add) |
        .vcpu_max_allocated_percent = (
            if .vcpu_allocatable == 0 then 0 else (
                .vcpu_max_allocated * 100 / .vcpu_allocatable * 10 | floor / 10
            ) end
        ) |
        .vcpu_available = 0 + ([.hyps[] | .vcpu_available] | add) |
        .vcpu_available_percent = (
            if .vcpu_allocatable == 0 then 0 else (
                .vcpu_available * 100 / .vcpu_allocatable * 10 | floor / 10
            ) end
        )
    ' <<< "$stats"
}
