#!/bin/sh
#
# Powered by Arviol
# Ilya Vasilyev, vStack, 2026
#

set -u

MODE="${1:-check}"

PD_BIN=/usr/local/sbin/vstack-vm-perfdata
PD_SVC=vstack_vm_perfdata
PD_OLD='/usr/sbin/bhyvectl --get-all --vm={}'
PD_NEW='/usr/local/sbin/vpd-bhyvectl --vm={}'
WRAPPER=/usr/local/sbin/vpd-bhyvectl
PL_BIN=/usr/local/bin/vStack-vm-perfdata-cpu-feeder
PL_SVC=vStackVMperfdataCPUfeeder
PL_OLD='bhyvectl --get-all --vm='
PL_NEW='bhyvectl --get-stats --get-cpu-topology --vm='
CONF=/usr/local/etc/vstack-node.conf
FREEZE_STAT='number of times idle requested at exit'

case "$MODE" in
    check|apply|revert) ;;
    *) echo "usage: $0 [check|apply|revert]" >&2; exit 1 ;;
esac

[ "$(id -u)" -eq 0 ] || { echo "ERROR: must run as root" >&2; exit 1; }

echo "== perfdata-fix on $(hostname), mode: $MODE =="

OSREL=$(uname -r)
OSMAJOR=${OSREL%%.*}
echo "os release:         $OSREL"

svc_running() {
    service "$1" status >/dev/null 2>&1
}

svc_restart() {
    service "$1" restart </dev/null >/dev/null 2>&1
    sleep 3
    svc_running "$1"
}

hits() {
    grep -obaF -- "$1" "$2" | wc -l | tr -d ' '
}

pd_state() {
    if [ ! -f "$PD_BIN" ]; then
        echo absent
        return
    fi
    n_old=$(hits "$PD_OLD" "$PD_BIN")
    n_new=$(hits "$PD_NEW" "$PD_BIN")
    if [ "$n_old" -eq 0 ] && [ "$n_new" -eq 0 ]; then
        echo not-affected
    elif [ "$n_old" -eq 0 ] && [ "$n_new" -eq 1 ] && [ -x "$WRAPPER" ]; then
        echo applied
    elif [ "$n_old" -eq 1 ] && [ "$n_new" -eq 0 ]; then
        echo needs-fix
    else
        echo unknown
    fi
}

pl_state() {
    if [ ! -f "$PL_BIN" ]; then
        echo absent
    elif grep -qF "$PL_OLD" "$PL_BIN"; then
        echo needs-fix
    elif grep -qF "$PL_NEW" "$PL_BIN"; then
        echo applied
    else
        echo unknown
    fi
}

interval() {
    awk -F= '
        /^[ \t]*\[/ { s = $0; gsub(/[ \t]/, "", s); next }
        $1 ~ /^[ \t]*interval[ \t]*$/ {
            v = $2; gsub(/[ \t\r]/, "", v)
            if (s == "[perfdata]") p = v
            if (s == "[influx]") i = v
        }
        END {
            r = (p != "") ? p : ((i != "") ? i : 10)
            if (r !~ /^[0-9]+$/ || r < 1) r = 10
            print r
        }' "$CONF" 2>/dev/null || echo 10
}

vm_list() {
    ls /dev/vmm 2>/dev/null | head -5
}

freezes() {
    total=0
    for vm in $VMS; do
        c=$(/usr/sbin/bhyvectl --get-stats --vm="$vm" 2>/dev/null |
            awk -F'\t' -v k="$FREEZE_STAT" 'index($1, k) == 1 { gsub(/ /, "", $2); print $2; exit }')
        total=$((total + ${c:-0}))
    done
    echo "$total"
}

measure() {
    VMS=$(vm_list)
    if [ -z "$VMS" ]; then
        echo "vCPU freezes:       no running VMs, not measured"
        FROZE=-1
        return
    fi
    nvm=$(echo "$VMS" | wc -l | tr -d ' ')
    win=$(( $(interval) * 2 + 5 ))
    [ "$win" -gt 125 ] && win=125
    echo "measuring vCPU0 freezes on $nvm VM(s) for $win s ..."
    a=$(freezes)
    sleep "$win"
    b=$(freezes)
    FROZE=$((b - a))
    echo "vCPU freezes:       $FROZE in $win s"
}

freeze_ok() {
    [ "$FROZE" -lt 0 ] && return 0
    if [ "$OSMAJOR" -ge 15 ]; then
        [ "$FROZE" -le $((20 * nvm)) ]
    else
        [ "$FROZE" -eq 0 ]
    fi
}

patch_pd() {
    off=$(grep -obaF -- "$PD_OLD" "$PD_BIN" | cut -d: -f1)
    case "$off" in
        ''|*[!0-9]*) echo "ERROR: cannot locate the bhyvectl call in $PD_BIN" >&2; return 1 ;;
    esac
    len=${#PD_OLD}
    if [ "$(dd if="$PD_BIN" bs=1 skip="$off" count="$len" 2>/dev/null)" != "$PD_OLD" ]; then
        echo "ERROR: unexpected bytes at offset $off in $PD_BIN" >&2
        return 1
    fi
    tmp="$PD_BIN.hotfix-tmp"
    cp -p "$PD_BIN" "$tmp" || return 1
    printf '%s' "$PD_NEW" | dd of="$tmp" bs=1 seek="$off" conv=notrunc 2>/dev/null
    diffs=$(cmp -l "$PD_BIN" "$tmp" | wc -l | tr -d ' ')
    if [ "$(hits "$PD_NEW" "$tmp")" -ne 1 ] || [ "$diffs" -lt 1 ] || [ "$diffs" -gt "$len" ]; then
        rm -f "$tmp"
        echo "ERROR: patched copy failed verification, $PD_BIN left as is" >&2
        return 1
    fi
    cp -p "$PD_BIN" "$PD_BIN.orig" || { rm -f "$tmp"; return 1; }
    mv "$tmp" "$PD_BIN" && PD_DONE=yes
}

write_wrapper() {
    printf '#!/bin/sh\nexec /usr/sbin/bhyvectl --get-stats --get-cpu-topology "$@"\n' > "$WRAPPER.tmp" &&
        chown root:wheel "$WRAPPER.tmp" && chmod 755 "$WRAPPER.tmp" && mv "$WRAPPER.tmp" "$WRAPPER"
}

patch_pl() {
    cp -p "$PL_BIN" "$PL_BIN.orig" || return 1
    tmp="$PL_BIN.hotfix-tmp"
    cp -p "$PL_BIN" "$tmp" || return 1
    sed -i '' "s|$PL_OLD|$PL_NEW|g" "$tmp"
    if grep -qF "$PL_OLD" "$tmp" || ! grep -qF "$PL_NEW" "$tmp"; then
        rm -f "$tmp"
        echo "ERROR: patched copy of $PL_BIN failed verification, left as is" >&2
        return 1
    fi
    mv "$tmp" "$PL_BIN" && PL_DONE=yes
}

restore() {
    [ -f "$1.orig" ] || return 0
    mv "$1.orig" "$1"
    echo "restored: $1"
}

restart_if() {
    [ "$2" = yes ] || return 0
    if svc_restart "$1"; then
        echo "service $1:  restarted, running"
    else
        echo "service $1:  FAILED to start" >&2
        return 1
    fi
}

PD=$(pd_state)
PL=$(pl_state)
PD_RUN=no; svc_running "$PD_SVC" && PD_RUN=yes
PL_RUN=no; svc_running "$PL_SVC" && PL_RUN=yes
echo "C++ collector:      $PD, running: $PD_RUN"
echo "Perl CPU feeder:    $PL, running: $PL_RUN"

if [ "$MODE" = revert ]; then
    [ "$PD" = applied ] && restore "$PD_BIN"
    [ "$PL" = applied ] && restore "$PL_BIN"
    [ "$(pd_state)" = applied ] || rm -f "$WRAPPER"
    rc=0
    restart_if "$PD_SVC" "$PD_RUN" || rc=1
    restart_if "$PL_SVC" "$PL_RUN" || rc=1
    echo "result: reverted"
    exit "$rc"
fi

if [ "$PD" = absent ] && [ "$PL" = absent ]; then
    echo "perfdata collector is not installed: fixes are not applicable, nothing changed"
    exit 0
fi

if [ "$PD" = unknown ] || [ "$PL" = unknown ]; then
    echo "ERROR: collector does not match any known build, nothing changed" >&2
    echo "send this output to support" >&2
    exit 1
fi

TODO=""
[ "$PD" = needs-fix ] && TODO="$TODO cxx-collector"
[ "$PL" = needs-fix ] && TODO="$TODO perl-feeder"

if [ -z "$TODO" ]; then
    measure
    if freeze_ok; then
        echo "verdict: everything is in order, nothing to do"
        exit 0
    fi
    echo "verdict: no collector needs the fix, but vCPU freezes go on: send this output to support"
    exit 2
fi
echo "verdict: fixes needed:$TODO"

if [ "$MODE" = check ]; then
    measure
    echo "run '$0 apply' to apply the fixes above"
    exit 3
fi

rc=0
PD_DONE=no
PL_DONE=no
if [ "$PD" = needs-fix ]; then
    write_wrapper && patch_pd || rc=1
    if [ "$rc" -eq 0 ]; then
        echo "applied: cxx-collector, backup: $PD_BIN.orig"
        restart_if "$PD_SVC" "$PD_RUN" || rc=1
    fi
fi
if [ "$rc" -eq 0 ] && [ "$PL" = needs-fix ]; then
    patch_pl || rc=1
    if [ "$rc" -eq 0 ]; then
        echo "applied: perl-feeder, backup: $PL_BIN.orig"
        restart_if "$PL_SVC" "$PL_RUN" || rc=1
    fi
fi

if [ "$rc" -ne 0 ]; then
    echo "result: apply failed, rolling back" >&2
    [ "$PD_DONE" = yes ] && restore "$PD_BIN"
    [ "$PL_DONE" = yes ] && restore "$PL_BIN"
    [ "$PD" = needs-fix ] && rm -f "$WRAPPER"
    rm -f "$PD_BIN.hotfix-tmp" "$PL_BIN.hotfix-tmp"
    restart_if "$PD_SVC" "$PD_RUN"
    restart_if "$PL_SVC" "$PL_RUN"
    exit 1
fi

measure
if freeze_ok; then
    echo "result: OK"
    exit 0
fi
echo "result: collector is fixed, but vCPU freezes go on: send this output to support"
exit 2