#!/bin/sh
# smoketest.sh — run on a live ISO via SSH to verify all critical components.
#
# Usage:
#   ssh root@<ip> 'sh -s' < smoketest.sh
#   or: scp smoketest.sh root@<ip>:/var/tmp/ && ssh root@<ip> sh /var/tmp/smoketest.sh
#
# Exit code: 0 = all required checks passed, 1 = at least one required check failed.

PASS=0
FAIL=0
WARN=0

ok()   { echo "[  OK  ] $*"; PASS=$((PASS+1)); }
fail() { echo "[ FAIL ] $*"; FAIL=$((FAIL+1)); }
warn() { echo "[ WARN ] $*"; WARN=$((WARN+1)); }
info() { echo "[ INFO ] $*"; }

echo "========================================"
echo " bee live ISO smoketest"
echo " host:   $(uname -n)"
echo " kernel: $(uname -r)"
echo " date:   $(date -u)"
echo "========================================"
echo ""

KVER=$(uname -r)
info "kernel: $KVER"
NVIDIA_BOOT_MODE="normal"
NVIDIA_MODULES_FLAVOR="proprietary"
for arg in $(cat /proc/cmdline 2>/dev/null); do
    case "$arg" in
        bee.nvidia.mode=*)
            NVIDIA_BOOT_MODE="${arg#*=}"
            ;;
    esac
done
if [ -f /etc/bee-nvidia-modules-flavor ]; then
    NVIDIA_MODULES_FLAVOR="$(tr -d '[:space:]' </etc/bee-nvidia-modules-flavor 2>/dev/null || echo proprietary)"
fi
info "nvidia boot mode: ${NVIDIA_BOOT_MODE}"
info "nvidia modules flavor: ${NVIDIA_MODULES_FLAVOR}"

# --- PATH & binaries ---
echo "-- PATH & binaries --"
for tool in dmidecode smartctl nvme ipmitool lspci bee; do
    if p=$(PATH="/usr/local/bin:/usr/sbin:/sbin:$PATH" command -v "$tool" 2>/dev/null); then
        ok "$tool found: $p"
    else
        fail "$tool: NOT FOUND"
    fi
done

if p=$(PATH="/usr/local/bin:$PATH" command -v nvidia-smi 2>/dev/null); then
    ok "nvidia-smi found: $p"
else
    fail "nvidia-smi: NOT FOUND"
fi

if p=$(PATH="/usr/local/bin:$PATH" command -v dcgmi 2>/dev/null); then
    ok "dcgmi found: $p"
else
    fail "dcgmi: NOT FOUND"
fi

if p=$(PATH="/usr/local/bin:$PATH" command -v nv-hostengine 2>/dev/null); then
    ok "nv-hostengine found: $p"
else
    fail "nv-hostengine: NOT FOUND"
fi

DCGM_PROFTESTER=""
for tool in dcgmproftester dcgmproftester13 dcgmproftester12 dcgmproftester11; do
    if p=$(PATH="/usr/local/bin:$PATH" command -v "$tool" 2>/dev/null); then
        DCGM_PROFTESTER="$p"
        break
    fi
done
if [ -n "$DCGM_PROFTESTER" ]; then
    ok "dcgmproftester found: $DCGM_PROFTESTER"
else
    fail "dcgmproftester: NOT FOUND"
fi

for tool in bee-gpu-burn bee-john-gpu-stress bee-nccl-gpu-stress all_reduce_perf; do
    if p=$(PATH="/usr/local/bin:$PATH" command -v "$tool" 2>/dev/null); then
        ok "$tool found: $p"
    else
        fail "$tool: NOT FOUND"
    fi
done

if p=$(PATH="/usr/local/bin:$PATH" command -v nvbandwidth 2>/dev/null); then
    ok "nvbandwidth found: $p"
else
    warn "nvbandwidth: NOT FOUND"
fi

echo ""
echo "-- NVIDIA modules --"
KO_DIR="/usr/local/lib/nvidia"
if [ -d "$KO_DIR" ]; then
    ko_count=$(ls "$KO_DIR"/*.ko 2>/dev/null | wc -l)
    ok "NVIDIA ko dir exists: $KO_DIR ($ko_count .ko files)"
else
    fail "NVIDIA ko dir missing: $KO_DIR"
fi

if /sbin/lsmod 2>/dev/null | grep -q "^nvidia "; then
    ok "module loaded: nvidia"
else
    fail "module NOT loaded: nvidia"
fi

for mod in nvidia_modeset nvidia_uvm; do
    if /sbin/lsmod 2>/dev/null | grep -q "^$mod "; then
        ok "module loaded: $mod"
    elif [ "${NVIDIA_MODULES_FLAVOR}" = "proprietary" ] && { [ "${NVIDIA_BOOT_MODE}" = "normal" ] || [ "${NVIDIA_BOOT_MODE}" = "full" ]; }; then
        fail "module NOT loaded in normal mode: $mod"
    elif [ "${NVIDIA_MODULES_FLAVOR}" = "proprietary" ]; then
        warn "module not loaded in GSP-off mode: $mod"
    else
        fail "module NOT loaded: $mod"
    fi
done

echo ""
echo "-- NVIDIA device nodes --"
for dev in nvidiactl nvidia0; do
    if [ -e "/dev/$dev" ]; then
        ok "/dev/$dev exists"
    else
        fail "/dev/$dev missing — nvidia-smi will return NVML_ERROR_LIBRARY_NOT_FOUND"
    fi
done

if [ -e /dev/nvidia-uvm ]; then
    ok "/dev/nvidia-uvm exists"
elif [ "${NVIDIA_MODULES_FLAVOR}" = "proprietary" ] && { [ "${NVIDIA_BOOT_MODE}" = "normal" ] || [ "${NVIDIA_BOOT_MODE}" = "full" ]; }; then
    fail "/dev/nvidia-uvm missing in normal mode"
elif [ "${NVIDIA_MODULES_FLAVOR}" = "proprietary" ]; then
    warn "/dev/nvidia-uvm missing — CUDA stress path may be unavailable until loaded on demand"
else
    fail "/dev/nvidia-uvm missing"
fi

echo ""
echo "-- nvidia-smi --"
if PATH="/usr/local/bin:$PATH" command -v nvidia-smi >/dev/null 2>&1; then
    if PATH="/usr/local/bin:$PATH" nvidia-smi -L 2>/dev/null | grep -q "GPU"; then
        gpu_count=$(PATH="/usr/local/bin:$PATH" nvidia-smi -L 2>/dev/null | grep -c "GPU")
        ok "nvidia-smi: $gpu_count GPU(s) found"
    else
        fail "nvidia-smi: runs but no GPUs detected"
    fi
else
    fail "nvidia-smi: not found in PATH"
fi

echo ""
echo "-- OpenCL / John --"
if [ -f /etc/OpenCL/vendors/nvidia.icd ]; then
    ok "OpenCL ICD present: /etc/OpenCL/vendors/nvidia.icd"
else
    fail "OpenCL ICD missing: /etc/OpenCL/vendors/nvidia.icd"
fi

if ldconfig -p 2>/dev/null | grep -q "libnvidia-opencl.so.1"; then
    ok "libnvidia-opencl.so.1 present in linker cache"
else
    fail "libnvidia-opencl.so.1 missing from linker cache"
fi

if command -v clinfo >/dev/null 2>&1; then
    if clinfo -l 2>/dev/null | grep -q "Platform"; then
        ok "clinfo: OpenCL platform detected"
    else
        fail "clinfo: no OpenCL platform detected"
    fi
else
    fail "clinfo: not found in PATH"
fi

if command -v john >/dev/null 2>&1; then
    if john --list=opencl-devices 2>/dev/null | grep -q "Device #"; then
        ok "john: OpenCL devices detected"
    else
        fail "john: no OpenCL devices detected"
    fi
else
    fail "john: not found in PATH"
fi

echo ""
echo "-- lib symlinks --"
for lib in libnvidia-ml libcuda; do
    if [ -f "/usr/lib/${lib}.so.1" ] || [ -L "/usr/lib/${lib}.so.1" ]; then
        ok "/usr/lib/${lib}.so.1 exists"
    else
        fail "/usr/lib/${lib}.so.1 MISSING (nvidia-smi will fail)"
    fi
done

echo ""
echo "-- systemd services --"
for svc in bee-nvidia bee-network bee-preflight bee-audit bee-web; do
    if systemctl is-active --quiet "$svc" 2>/dev/null; then
        ok "service active: $svc"
    else
        fail "service NOT active: $svc"
    fi
done

if systemctl is-active --quiet bee-selfheal.timer 2>/dev/null; then
    ok "timer active: bee-selfheal.timer"
else
    fail "timer NOT active: bee-selfheal.timer"
fi

echo ""
echo "-- runtime health --"
if [ -f /appdata/bee/export/runtime-health.json ] && [ -s /appdata/bee/export/runtime-health.json ]; then
    ok "runtime: runtime-health.json present and non-empty"
else
    fail "runtime: runtime-health.json missing or empty"
fi

if [ -f /appdata/bee/export/runtime-health.log ]; then
    info "last runtime log line: $(tail -1 /appdata/bee/export/runtime-health.log)"
else
    warn "runtime: no log found at /appdata/bee/export/runtime-health.log"
fi

for svc in ssh bee-sshsetup; do
    if systemctl is-active --quiet "$svc" 2>/dev/null \
        || systemctl show "$svc" --property=ActiveState 2>/dev/null | grep -q "inactive\|exited"; then
        ok "service ok: $svc"
    else
        warn "service status unknown: $svc"
    fi
done

echo ""
echo "-- bee binary --"
BEE=/usr/local/bin/bee
if [ -x "$BEE" ]; then
    ok "bee binary: present"
    ver=$("$BEE" version 2>/dev/null || echo "unknown")
    info "bee version: $ver"
else
    fail "bee binary: NOT FOUND at $BEE"
fi

echo ""
echo "-- audit last run --"
if [ -f /appdata/bee/export/bee-audit.json ] && [ -s /appdata/bee/export/bee-audit.json ]; then
    ok "audit: bee-audit.json present and non-empty"
    info "size: $(du -sh /appdata/bee/export/bee-audit.json | cut -f1)"
else
    fail "audit: bee-audit.json missing or empty"
fi

if [ -f /appdata/bee/export/bee-audit.log ]; then
    last_line=$(tail -1 /appdata/bee/export/bee-audit.log)
    info "last log line: $last_line"
    if grep -q "audit output written" /appdata/bee/export/bee-audit.log 2>/dev/null; then
        ok "audit: completed successfully"
    else
        warn "audit: 'audit output written' not found in log — may have failed"
    fi
    if grep -q "nvidia: enrichment skipped\|nvidia.*skipped\|enrichment skipped" /appdata/bee/export/bee-audit.log 2>/dev/null; then
        reason=$(grep -E "nvidia.*skipped|enrichment skipped" /appdata/bee/export/bee-audit.log | tail -1)
        fail "audit: nvidia enrichment skipped — $reason"
    else
        ok "audit: nvidia enrichment OK (no skip message)"
    fi
else
    warn "audit: no log found at /appdata/bee/export/bee-audit.log"
fi

echo ""
echo "-- bee web --"
if [ -f /appdata/bee/export/bee-web.log ]; then
    info "last web log line: $(tail -1 /appdata/bee/export/bee-web.log)"
else
    warn "web: no log found at /appdata/bee/export/bee-web.log"
fi

if bash -c 'exec 3<>/dev/tcp/127.0.0.1/80 && printf "GET /healthz HTTP/1.0\r\nHost: localhost\r\n\r\n" >&3 && grep -q "^ok$" <&3'; then
    ok "web: health endpoint reachable on 127.0.0.1:80"
else
    fail "web: health endpoint not reachable on 127.0.0.1:80"
fi

echo ""
echo "-- network --"
if ip route show default 2>/dev/null | grep -q "default"; then
    gw=$(ip route show default | awk '{print $3}' | head -1)
    ok "default route: $gw"
else
    fail "no default route"
fi
if ping -c1 -W3 1.1.1.1 >/dev/null 2>&1; then
    ok "internet: reachable (1.1.1.1)"
else
    warn "internet: unreachable (expected — live CD runs in isolated network segment)"
fi

echo ""
echo "-- /etc/profile.d/bee.sh PATH --"
if grep -q "/usr/local/bin" /etc/profile.d/bee.sh 2>/dev/null; then
    ok "/etc/profile.d/bee.sh exports /usr/local/bin"
else
    fail "/etc/profile.d/bee.sh does not add /usr/local/bin to PATH"
fi

echo ""
echo "========================================"
echo " Results: OK=$PASS  FAIL=$FAIL  WARN=$WARN"
echo "========================================"
[ "$FAIL" -eq 0 ] && exit 0 || exit 1
