#!/bin/sh
# Copyright (c) 2026 Jiaxing Hu <gahing@gahingwoo.com>
# SPDX-License-Identifier: GPL-2.0
#
# charsiu-doctor: is this board able to run charsiu on the NPU, and if not,
# which part is missing? Prints a plain report and exits non-zero if anything
# CRITICAL failed, so it is scriptable and is the thing to paste into an issue.
#
#   charsiu-doctor        the report
#   charsiu-doctor -q     failures and the verdict only
#   charsiu-doctor --run  also generate a few tokens (slow: staging is ~20 s)
set -u

CHARSIU_HERE=$(dirname "$0")
for _l in ${CHARSIU_LIB:+"$CHARSIU_LIB/charsiu-lib.sh"} \
	  "$CHARSIU_HERE/charsiu-lib.sh" \
	  "$CHARSIU_HERE/../scripts/charsiu-lib.sh" \
	  "$CHARSIU_HERE/../../opt/charsiu/charsiu-lib.sh" \
	  /opt/charsiu/charsiu-lib.sh; do
	[ -r "$_l" ] && { . "$_l"; break; }
done
command -v ini_get >/dev/null 2>&1 \
	|| { echo "$(basename "$0"): charsiu-lib.sh not found" >&2; exit 1; }

CONF=$(charsiu_conf)
QUIET=0; DORUN=0; PASTE=0
for a in "$@"; do
	case "$a" in
	-q) QUIET=1 ;;
	--run) DORUN=1 ;;
	# --paste is for an issue report, so it must be plain: no colour, no
	# ANSI, and it carries the things that are always asked for anyway --
	# which kernel, which dtb, which model, what rate.
	--paste) PASTE=1 ;;
	esac
done
if [ "$PASTE" = 1 ]; then
	printf '```\n'
	printf 'charsiu-doctor --paste\n'
	printf 'kernel   %s\n' "$(uname -r)"
	# tr succeeds on empty input, so `cat ... | tr ... || uname -m` never
	# reaches the fallback. Test the file instead of the pipeline.
	if [ -r /proc/device-tree/model ]; then
		printf 'machine  %s\n' "$(tr -d '\0' < /proc/device-tree/model)"
	else
		printf 'machine  %s\n' "$(uname -m)"
	fi
	printf 'accel    %s\n' "$([ -e "${CHARSIU_ACCEL_DEV:-/dev/accel/accel0}" ] && echo present || echo MISSING)"
	n=0; for d in /sys/bus/platform/devices/*.npu; do [ -d "$d" ] && n=$((n+1)); done
	printf 'npu      %d core(s)\n' "$n"
	for r in /sys/class/regulator/regulator.*; do
		[ -r "$r/name" ] || continue
		[ "$(cat "$r/name")" = vdd_npu_s0 ] || continue
		printf 'rail     vdd_npu_s0 %s\n' "$(cat "$r/state" 2>/dev/null)"
	done
	# grep -c EXITS 1 WHEN THE COUNT IS ZERO, so `|| echo ?` fires on the
	# perfectly good answer "0" and both get printed. Capture, then default.
	nf=$(dmesg 2>/dev/null | grep -c 'MMU_DTE_ADDR\|job timed out') || nf=${nf:-0}
	printf 'faults   %s in dmesg\n' "${nf:-?}"
	printf '\n'
	CHARSIU_CONFIG="$CONF" "$0" 2>&1 | sed 's/\x1b\[[0-9;]*m//g'
	printf '```\n'
	exit 0
fi

if [ -t 1 ]; then
	C_G=$(printf '\033[1;32m'); C_R=$(printf '\033[1;31m'); C_Y=$(printf '\033[1;33m')
	C_D=$(printf '\033[2m'); C_B=$(printf '\033[1m'); C_0=$(printf '\033[0m')
else C_G=; C_R=; C_Y=; C_D=; C_B=; C_0=; fi

CRIT=0; WARN=0; OKC=0
ok()   { OKC=$((OKC+1)); [ "$QUIET" = 1 ] || printf '  %s[ OK ]%s %s\n' "$C_G" "$C_0" "$*"; }
bad()  { printf '  %s[FAIL]%s %s\n' "$C_R" "$C_0" "$*"; CRIT=$((CRIT+1)); }
warn() { printf '  %s[WARN]%s %s\n' "$C_Y" "$C_0" "$*"; WARN=$((WARN+1)); }
info() { [ "$QUIET" = 1 ] || printf '  %s[INFO]%s %s\n' "$C_D" "$C_0" "$*"; }
hdr()  { [ "$QUIET" = 1 ] || printf '\n%s%s%s\n' "$C_B" "$*" "$C_0"; }


hdr "the hardware"
# The mainline rocket driver is an accel device, not a render node.
# CHARSIU_ACCEL_DEV names it, as it does for charsiu-install.sh. A board with
# more than one accel device should not need the two tools to disagree.
ACCEL="${CHARSIU_ACCEL_DEV:-/dev/accel/accel0}"
if [ -e "$ACCEL" ]; then
	ok "$ACCEL is present"
	[ -r "$ACCEL" ] && [ -w "$ACCEL" ] \
		&& ok "and this user can open it" \
		|| bad "but this user cannot open it (group/permissions on $ACCEL)"
else
	bad "$ACCEL is missing, so the rocket driver did not bind"
fi

ncore=0
for d in /sys/bus/platform/devices/27700000.npu /sys/bus/platform/devices/27708000.npu; do
	[ -d "$d" ] || continue
	ncore=$((ncore + 1))
	[ -e "$d/driver" ] || warn "${d##*/} exists but has no driver bound"
done
case "$ncore" in
2) ok "both NPU cores are present (27700000.npu, 27708000.npu)" ;;
1) warn "only ONE NPU core is enabled. charsiu's int4 decode wants both, and one core is about 1.4x slower" ;;
0) bad "no NPU core in the device tree" ;;
esac

# Informational: with the rail handed to the power domain it reads disabled at
# rest, and that is CORRECT. It does not mean the NPU is broken.
for r in /sys/class/regulator/regulator.*; do
	[ -r "$r/name" ] || continue
	[ "$(cat "$r/name")" = "vdd_npu_s0" ] || continue
	st=$(cat "$r/state" 2>/dev/null)
	info "vdd_npu_s0 is $st at rest (disabled here is correct: the rail follows the power domain)"
done

#
# WHERE THE THREADS RUN, because on a big.LITTLE part it is worth 27.7% and
# until 2026-09-11 nothing here mentioned it at all.
#
# Measured on a ROCK 4D, Qwen3-0.6B, ten passes: unpinned decode is BIMODAL
# across a 28% spread -- 20.5 or 26.4 depending on where the scheduler put four
# threads -- and pinned to the four A72s it is 26.4 with a 1.6% spread. charsiu
# pins its calling thread to the fastest cluster itself now; this says whether
# that will happen on THIS machine, because a default that silently does
# nothing is the thing this tree keeps being caught by.
#
CL=$(for c in /sys/devices/system/cpu/cpu[0-9]*/cpufreq/cpuinfo_max_freq; do
	[ -r "$c" ] && cat "$c"; done | sort -u | wc -l)
if [ "${CL:-0}" -eq 0 ]; then
	info "no cpufreq here, so charsiu will not pin anything (it needs cpuinfo_max_freq to tell the clusters apart)"
elif [ "$CL" -eq 1 ]; then
	info "one CPU cluster: nothing to pin to, and charsiu leaves the scheduler alone"
else
	FAST=$(for c in /sys/devices/system/cpu/cpu[0-9]*; do
		f=$(cat "$c/cpufreq/cpuinfo_max_freq" 2>/dev/null || echo 0)
		echo "$f ${c##*cpu}"; done | sort -rn |
		awk 'NR==1{t=$1} $1==t{printf "%s%s", (n++?",":""), $2}')
	TOP=$(for c in /sys/devices/system/cpu/cpu[0-9]*/cpufreq/cpuinfo_max_freq; do
		[ -r "$c" ] && cat "$c"; done | sort -rn | head -1)
	ok "$CL CPU clusters; charsiu pins its calling thread to the fastest, CPUs $FAST at $((TOP / 1000)) MHz"
	info "CHARSIU_AFFINITY=0 leaves it to the scheduler, which on this part costs a bimodal decode"
	G=$(cat /sys/devices/system/cpu/cpu0/cpufreq/scaling_governor 2>/dev/null)
	[ "$G" = performance ] ||
		info "the governor is ${G:-unknown}, not performance -- every timing here will be noisier than it needs to be"
fi

hdr "the configuration"
if [ -f "$CONF" ]; then
	ok "config at $CONF"
	# SAY WHICH CHANNEL. A dev install has the probes on it and a paste
	# from one is a different thing to read: npu_gemm_test and
	# charsiu_matmul can wedge the block, so "it hung" from a dev board and
	# from a stable one are not the same report.
	_ch=$(ini_get update channel); : "${_ch:=stable}"
	case "$_ch" in
	dev) info "channel dev: the hardware probes are installed" ;;
	*)   ok "channel $_ch" ;;
	esac
else
	bad "no config at $CONF. Run charsiu-config"
fi
KS=$(ini_get npu kslice)
[ -n "$KS" ] && info "K slice $KS (written to both CHARSIU_NPU_KMAX and CHARSIU_NPU_W4_GROUP; they must match)"

hdr "the model"
M=$(ini_get model path)
CK=$(find_bin charsiu_check)
if [ -z "$M" ]; then
	bad "[model] path is not set"
elif [ ! -f "$M" ]; then
	bad "$M does not exist. charsiu-get fetches one"
else
	# busybox's stat needs FEATURE_STAT_FORMAT for -c and a rootfs here was
	# built without it, so a 1321 MB model reported as 0 MB and the memory
	# check below became meaningless. That is why file_bytes falls back to
	# ls -ln rather than to wc -c: busybox's wc READS THE WHOLE FILE.
	sz=$(( $(file_bytes "$M") / 1048576 ))
	ok "$M (${sz} MB)"
	if [ -n "$CK" ]; then
		# EXIT 3 IS NOT A FAILURE, it is "runnable, but not as a chat
		# model". Reported as CANNOT it sends somebody to re-download a
		# file that is fine and is simply in the wrong setting.
		out=$("$CK" -q "$M" 2>&1) && _r=0 || _r=$?
		case "$_r" in
		0) ok "charsiu can run it: $out" ;;
		3) bad "[model] path points at something that is not a chat model:
     $out
     charsiu-config, or charsiu pull, to set a chat model" ;;
		*) bad "charsiu CANNOT run it: $out" ;;
		esac
	else
		warn "charsiu_check not found, so the file was not verified"
	fi
	# Staging builds a second, requantised copy of every routed tensor, so
	# peak RSS runs well above the file size. Measured: a 1321 MB q8_0 file
	# peaks at about 2535 MB.
	free=$(awk '/^MemAvailable:/ {print int($2/1024)}' /proc/meminfo 2>/dev/null)
	if [ -n "$free" ]; then
		need=$((sz * 2))
		if [ "$free" -lt "$need" ]; then
			warn "${free} MB available and staging this model peaks near ${need} MB"
		else
			ok "${free} MB available, staging peaks near ${need} MB"
		fi
	fi
fi

# WHAT ELSE IS ON THIS MACHINE, because the models directory stopped holding
# only chat models and a person who pulled a whisper model has no way to ask
# whether the binary that runs it got installed.
hdr "the other things it can run"
_any=0
for _f in $(all_models 2>/dev/null) "$(ini_get model path)"; do
	[ -f "$_f" ] || continue
	[ -n "$CK" ] || continue
	_o=$("$CK" -q "$_f" 2>&1) && continue     # 0: an ordinary chat model
	case "$_o" in
	VISION-TOWER*) _any=1; info "$(basename "$_f"): a vision tower" ;;
	WHISPER*)      _any=1; info "$(basename "$_f"): a whisper model" ;;
	esac
done
for _b in charsiu_vision charsiu_clip charsiu_whisper; do
	if [ -n "$(find_bin $_b)" ]; then
		ok "$_b is installed"
	else
		warn "$_b is NOT installed, so a picture or a clip has nothing to run it"
	fi
done
[ "$_any" = 1 ] || info "no vision towers or whisper models in the models directory"

hdr "the kernel's opinion"
if command -v dmesg >/dev/null 2>&1; then
	n=$(dmesg 2>/dev/null | grep -c 'MMU_DTE_ADDR\|job timed out' || true)
	if [ "${n:-0}" -gt 0 ]; then
		warn "$n iommu-fault or job-timeout lines in dmesg"
	else
		ok "no iommu faults and no job timeouts in dmesg"
	fi
else
	info "dmesg unavailable, so the kernel log was not read"
fi

if [ "$DORUN" = 1 ] && [ "$CRIT" = 0 ]; then
	hdr "an actual decode"
	L=$(find_bin charsiu); [ -n "$L" ] || L=$(dirname "$0")/charsiu
	info "staging takes about 20 s before the first token"
	if o=$("$L" -p "The capital of France is" -n 8 -q 2>&1); then
		r=$(echo "$o" | grep -oE 'gen [0-9]+ tok in [0-9]+ ms, [0-9.]+ tok/s' || true)
		[ -n "$r" ] && ok "$r" || warn "it ran but printed no rate"
	else
		bad "the decode failed"
	fi
fi

[ "$QUIET" = 1 ] || printf '\n'
if [ "$CRIT" -gt 0 ]; then
	printf '%sVERDICT%s %d passed, %d warning(s), %s%d CRITICAL%s\n' \
		"$C_B" "$C_0" "$OKC" "$WARN" "$C_R" "$CRIT" "$C_0"
	exit 1
fi
printf '%sVERDICT%s %d passed, %d warning(s), nothing critical\n' \
	"$C_B" "$C_0" "$OKC" "$WARN"
exit 0
