#!/bin/sh
# Copyright (c) 2026 Jiaxing Hu <gahing@gahingwoo.com>
# SPDX-License-Identifier: GPL-2.0
#
# charsiu: run the model named in config.ini -- $HOME/.charsiu's if there is
# one, /etc/charsiu's otherwise.
#
# The runtime reads 109 CHARSIU_* environment variables. Six of them decide
# whether a decode goes to the NPU at all, and getting one wrong does not fail
# loudly. It quietly falls back to the CPU and runs a third as fast. Nobody
# should have to remember this:
#
#   CHARSIU_NPU=1 CHARSIU_NPU_QUANT=1 CHARSIU_NPU_W4V=1 CHARSIU_NPU_KMAX=1024 \
#   CHARSIU_NPU_W4_GROUP=1024 CHARSIU_NPU_MAXN=262144 CHARSIU_COEF_ELEMS=65536 \
#   charsiu_run /opt/charsiu/models/....gguf -p "..." -n 64 -c 512 -t 4
#
# KMAX AND GROUP MUST BE EQUAL. npudev.c gates the int4 path on
# `t->kgroup == g->kmax`; set them differently and the tensor silently takes
# the other path. config.ini keeps one value and writes both.
#
# Usage:  charsiu                    a conversation (the model stays staged)
#         charsiu -p PROMPT [-n N]   one answer and out
#         charsiu bench              what this board does, in tok/s
#         charsiu --show             print the command it would run, and stop
#         charsiu --config FILE      use a different config
set -eu

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)
RUN="${CHARSIU_RUN_BIN:-}"
SHOW=0

BENCH=0
while [ $# -gt 0 ]; do
	case "$1" in
	--config) CONF="$2"; shift 2 ;;
	--show)   SHOW=1; shift ;;
	bench)    BENCH=1; shift ;;
	*)        break ;;
	esac
done

# A BARE WORD IS THE PROMPT, and it was not. `charsiu "what is a bandgap"`
# is line two of --help and the second line of the front door's own header,
# and the words went to charsiu_run as a positional it does not accept, so the
# most prominently documented form of this command answered with charsiu_run's
# usage. `charsiu run MODEL PROMPT` was broken the same way.
#
# Flags and their values pass through untouched, so `charsiu -n 32 "hello"`
# still means what it looks like and `charsiu -n 32` does not treat 32 as a
# question. The list of value-taking flags is charsiu_run's own.
_skip=0; _hasp=0; _bare=0
for a in "$@"; do
	if [ "$_skip" = 1 ]; then _skip=0; continue; fi
	case "$a" in
	-p|-f) _hasp=1; _skip=1 ;;
	--sys|-n|-c|-t|--temp|--top-p|--seed|--hold-secs|--cache-at|--image|--mmproj) _skip=1 ;;
	-*) ;;
	*) _bare=1 ;;
	esac
done
if [ "$_bare" = 1 ] && [ "$_hasp" = 0 ]; then
	_skip=0; _p=""; _n=$#
	while [ "$_n" -gt 0 ]; do
		a=$1; shift; _n=$((_n - 1))
		if [ "$_skip" = 1 ]; then set -- "$@" "$a"; _skip=0; continue; fi
		case "$a" in
		--sys|-n|-c|-t|--temp|--top-p|--seed|--hold-secs|--cache-at|--image|--mmproj)
			_skip=1; set -- "$@" "$a" ;;
		-*) set -- "$@" "$a" ;;
		*)  if [ -z "$_p" ]; then _p="$a"; else set -- "$@" "$a"; fi ;;
		esac
	done
	if [ -n "$_p" ]; then
		set -- "$@" -p "$_p"
		# A SINGLE BARE WORD IS AMBIGUOUS, AND GETTING IT WRONG WAS
		# SILENT. On an install predating a subcommand, `charsiu update`
		# was read as a question and the model answered with a cheerful
		# paragraph about workouts, which looks like the update ran and
		# printed something odd. One line on stderr says which way it
		# was read; a word with a space in it is plainly a question and
		# says nothing.
		case "$_p" in
		*[!A-Za-z0-9_-]*) ;;
		*) echo "charsiu: reading \"$_p\" as a question, not a command." \
			"\`charsiu --help\` lists the commands." >&2 ;;
		esac
		# ASKING IS NOT COMPLETING. --help draws exactly this line:
		# `charsiu "what is a bandgap"` asks once, `charsiu -p "..."` is
		# print mode for a pipe. Without the chat template the question was
		# echoed back and continued as a document -- "Name one color." came
		# back as "Name one color.\n3. A bold and bolder design...".
		# An explicit -p is left alone: that form chose completion.
		case " $* " in *" --chat "*|*" -i "*|*" --interactive "*) ;;
		*) set -- "$@" --chat ;; esac
		# AND -q, OR THE QUESTION COMES BACK WITH THE ANSWER. Without it
		# charsiu_run echoes the prompt it was given, which for a chat turn
		# is the whole template: somebody who asked "Name one color." got
		# "<|im_start|>system ... <|im_end|> ..." and then the answer.
		case " $* " in *" -q "*) ;; *) set -- "$@" -q ;; esac
	fi
fi

# NO ARGUMENTS MEANS A CONVERSATION, not the built-in demo prompt. Staging
# costs about twenty seconds and a one-shot pays it for a single answer, so the
# default has to be the mode that pays it once.
# NO PROMPT MEANS A CONVERSATION, and the test used to be "no arguments
# at all". `charsiu -n 10` therefore ran the built-in demo prompt instead
# of talking, which is not what anyone typing that meant. Reuse the scan
# above: a bare word or -p/-f is a question, anything else is not.
INTERACTIVE=0
_askish=0
case " $* " in *" --info "*|*" --tokens "*|*" --logits "*) _askish=1 ;; esac
[ "$_bare" = 0 ] && [ "$_hasp" = 0 ] && [ "$_askish" = 0 ] \
	&& [ "$BENCH" = 0 ] && [ "$SHOW" = 0 ] && INTERACTIVE=1

[ -f "$CONF" ] || { echo "charsiu: no config at $CONF (run charsiu-config)" >&2; exit 1; }

# `charsiu run MODEL ...` names it outright; otherwise the config decides.
MODEL="${CHARSIU_MODEL:-$(ini_get model path)}"
[ -n "$MODEL" ] && MODEL=$(resolve_model "$MODEL")

# THE SHIPPED CONFIG NAMES A MODEL NOBODY HAS YET. A clean install where the
# wizard downloaded SmolLM2 still ended at "Llama-3.2-1B-Instruct-Q8_0.gguf does
# not exist", because the default path is a placeholder and the download never
# wrote over it. What is on disk beats what the config guessed.
if [ -z "$MODEL" ]; then
	_n=$(all_models | wc -l)
	if [ "$_n" -eq 1 ]; then
		MODEL=$(all_models | head -1)
		[ "$SHOW" = 0 ] && echo "charsiu: using the only model here, ${MODEL##*/}" >&2
	elif [ "$_n" -gt 1 ]; then
		echo "charsiu: the configured model is not here. $_n models are:" >&2
		all_models | while IFS= read -r f; do echo "  ${f##*/}" >&2; done
		echo "  pick one with:  charsiu run NAME     or set it with:  charsiu-config" >&2
		exit 1
	fi
fi
[ -n "$MODEL" ] || { echo "charsiu: [model] path is not set in $CONF" >&2; exit 1; }
[ -f "$MODEL" ] || { echo "charsiu: $MODEL does not exist (charsiu pull fetches one)" >&2; exit 1; }

if [ -z "$RUN" ]; then
	RUN=$(find_bin charsiu_run)
fi
[ -n "$RUN" ] || { echo "charsiu: charsiu_run not found" >&2; exit 1; }

THREADS=$(ini_get run threads);  : "${THREADS:=4}"
CTX=$(ini_get run context);      : "${CTX:=512}"
CPUS=$(ini_get run cpus)
NPU=$(ini_get npu enable);       : "${NPU:=yes}"
KSLICE=$(ini_get npu kslice);    : "${KSLICE:=1024}"
MAXN=$(ini_get npu maxn);        : "${MAXN:=262144}"
COEF=$(ini_get npu coef_elems);  : "${COEF:=65536}"
GFIT=$(ini_get npu group_fit);   : "${GFIT:=no}"

set -- "$@"
E=""
INHERITED=""

# See env_default in charsiu-lib.sh: the config file is the default and an
# exported CHARSIU_* is the override, which is why these are not plain
# assignments into $E.
case "$NPU" in
yes|1|true|on)
	env_default CHARSIU_NPU 1
	env_default CHARSIU_NPU_QUANT 1
	env_default CHARSIU_NPU_W4V 1
	# kslice writes BOTH, and they must match. See the note above.
	env_default CHARSIU_NPU_KMAX "$KSLICE"
	env_default CHARSIU_NPU_W4_GROUP "$KSLICE"
	env_default CHARSIU_NPU_MAXN "$MAXN"
	env_default CHARSIU_COEF_ELEMS "$COEF"
	case "$GFIT" in
	yes|1|true|on) env_default CHARSIU_NPU_W4_GROUP_FIT 1 ;;
	esac ;;
esac
[ -n "$CPUS" ] && env_default CHARSIU_CPUS "$CPUS"
# THE POOL WANTS DIFFERENT CORES FROM THE CALLING THREAD, and the board
# measured both halves: a prompt is 4.3% faster with the pool on all eight,
# and decode is 26% SLOWER two runs in three if the CALLING thread can land on
# an A53. cpus keeps the caller on the fast cores, pool_cpus opens the machine
# to the workers. Neither has a default.
POOLCPUS=$(ini_get run pool_cpus)
[ -n "$POOLCPUS" ] && env_default CHARSIU_POOL_CPUS "$POOLCPUS"

# ONE SWITCH, NOT A SET OF ENVIRONMENT VARIABLES NOBODY WILL REMEMBER. The
# stage timings existed only as CHARSIU_STAGES, which meant they were reachable
# by whoever already knew the name. [debug] enable makes them a setting.
DEBUG=$(ini_get debug enable)
case "$DEBUG" in
yes|1|true|on) env_default CHARSIU_STAGES 1 ;;
esac

# anything under [extra] is passed through verbatim, which is the escape hatch
# for the other hundred variables without putting them in the TUI.
# A line that STARTS with # is a comment too. An earlier version only
# stripped ` #...` and so passed the section's own comment to env as an
# argument. Anything that is not literally CHARSIU_NAME=VALUE is dropped.
EX=$(awk '/^[ \t]*\[extra\][ \t]*$/{f=1;next} /^[ \t]*\[/{f=0}
	f {
		sub(/^[ \t]+/,""); sub(/[ \t]+#.*$/,"");
		if ($0 ~ /^#/) next;
		gsub(/[ \t]/,"");
		if ($0 ~ /^CHARSIU_[A-Za-z0-9_]+=/) printf "%s ", $0;
		else if (length($0)) printf "" > "/dev/stderr";
	}' "$CONF")
[ -n "$EX" ] && E="$E $EX"

# TWENTY SECONDS OF SILENCE IS THE MOST OFF-PUTTING PART OF A FIRST RUN, and
# the runtime already says enough to draw a bar: it announces how many tensors
# it will stage (n_layer * 7 + 1, seven projections a layer plus the output
# head, 113 for the model every board round uses, which is exactly what its logs
# say) and then heartbeats every sixteen. Progress goes to stderr, the generated
# text to stdout, so the two never mix.
# mawk batches its output and its fflush() does not help; -Winteractive does.
# Probe for the flag rather than sniffing which awk this is, and fall back
# quietly for the awks that reject it.
CHARSIU_AWK=awk
if awk -Winteractive '' </dev/null >/dev/null 2>&1; then
	CHARSIU_AWK="awk -Winteractive"
fi

stage_filter() {
	# DO NOT WRITE TO "/dev/stderr" FROM awk. mawk, which is Debian's awk,
	# opens it as an ordinary file and buffers it whole, and fflush() on that
	# name does not help. Measured: three lines fed a second apart all came
	# out together at the end. Meanwhile stdout goes straight to the terminal
	# with no filter in the way, so the ANSWER printed before the staging
	# messages that must precede it -- the board log has the bar and the
	# whole NPU report arriving after the generation summary.
	#
	# Write to plain stdout, flush every line, and let the caller point that
	# at stderr.
	#
	# AND fflush() IS NOT ENOUGH, BECAUSE mawk IGNORES IT. Measured on
	# mawk 1.3.4: two lines fed a second apart came out 1 ms apart with
	# `{ print; fflush() }`, and a second apart with `-Winteractive`. So the
	# flag is what does the work; the fflush() calls stay for awks that
	# honour them and have no such flag.
	$CHARSIU_AWK '''
	/staging [0-9]+$/ {
		match($0, /staging [0-9]+$/); tot = substr($0, RSTART + 8) + 0
		printf "  staging %d NPU tensors\n", tot; fflush(); next
	}
	/tensors staged/ {
		n = $3 + 0
		if (tot > 0) {
			pct = int(n * 100 / tot); if (pct > 100) pct = 100
			# THE HEARTBEAT ONLY FIRES EVERY SIXTEENTH TENSOR, so for
			# 113 the last one is 112 and the bar stops at 99% with no
			# later line to close it. The generated text then lands on
			# the same line: "99%  112/113OK, Bob!". Close it on the
			# last heartbeat there is going to be.
			if (n + 16 > tot) {
				printf "\r  [########################] 100%%  %d/%d\n",
				       tot, tot
				bar = 0
			} else {
				printf "\r  [%-24s] %3d%%  %d/%d",
				       substr("########################", 1, int(pct/4.2)),
				       pct, n, tot
				bar = 1
			}
			fflush()
		}
		next
	}
	# THE DENOMINATOR IS A PREDICTION AND THE BAR CLOSES ITSELF, so the
	# bar cannot notice a tensor that never staged. The heartbeat fires
	# every sixteenth tensor, the last one for 183 lands at 176, and the
	# rule above prints 183/183 to keep the answer off the bar line. A
	# gemma3 round drew a full bar while 182 of 183 had staged, and the one
	# missing was the output head -- 44% of that token, on the CPU, with a
	# 100% bar over it. The runtime says the real count when it is done.
	/^charsiu NPU: [0-9]+ tensors,/ {
		if (bar) { printf "\n"; bar = 0 }
		print; fflush()
		if (tot > 0 && $3 + 0 < tot)
			printf "  %d of the %d it meant to stage stayed on the CPU\n",
			       tot - ($3 + 0), tot
		fflush(); next
	}
	# a pass-through line lands mid-bar unless the bar is closed first,
	# and \r leaves the cursor at column 0 of a line that still has a bar on it.
	{ if (bar) { printf "\n"; bar = 0 }
	  print; fflush() }
	END { if (bar) printf "\r  [########################] 100%%  %d/%d\n", tot, tot; fflush() }
	'''
}

if [ "$BENCH" = 1 ]; then
	printf '\n  charsiu bench: %s\n' "$(basename "$MODEL")"
	printf '  %s, %s threads, context %s\n\n' \
	       "$(case "$NPU" in yes|1|true|on) echo "NPU int4";; *) echo "CPU";; esac)" \
	       "$THREADS" "$CTX"
	# 64 tokens is what every board round measures, so the number here is
	# comparable with the logs rather than being a new metric nobody can
	# check. --ignore-eos so a short answer does not shorten the sample.
	out=$( { env $E "$RUN" "$MODEL" -p "The capital of France is" -n 64 \
		--ignore-eos -q -c "$CTX" -t "$THREADS" 2>&1 1>&3 | stage_filter >&2; } 3>&1 )
	echo "$out" | grep -oE 'gen 64 tok in [0-9]+ ms, [0-9.]+ tok/s' \
		| sed 's/^/  /' || echo "  (no rate reported)"
	echo "$out" | grep -oE '[0-9.]+ GB/s of weights' | sed 's/^/  /' || true

	# THE ONE NUMBER THAT DECIDES THE NEXT PIECE OF WORK, and it differs
	# between a development host and the board. Batching the projections
	# measured 1.1x to 1.4x on a host with cache and memory to spare; the
	# board is 10 GB/s of LPDDR5 behind four A72s, where the same kernel
	# sits much closer to the memory roof. Only under debug, because it
	# walks every layer twice and takes a while.
	case "$DEBUG" in
	yes|1|true|on)
		BB=$(find_bin bench_batch)
		if [ -n "$BB" ]; then
			printf '\n  batching the projections, M=8\n'
			"$BB" "$MODEL" 8 2>/dev/null | grep -E "at a time|batched" \
				|| echo "  (bench_batch would not run)"
		else
			printf '\n  bench_batch is not installed\n'
		fi ;;
	esac
	printf '\n'
	exit 0
fi

if [ "$SHOW" = 1 ]; then
	printf 'env%s \\\n  %s %s -c %s -t %s' "$E" "$RUN" "$MODEL" "$CTX" "$THREADS"
	[ $# -gt 0 ] && printf ' %s' "$@"
	printf '\n'
	# what the environment is already deciding, which the line above cannot
	# show because those names are not being passed at all.
	[ -n "$INHERITED" ] && printf '# from the environment, not the config:%s\n' "$INHERITED"
	exit 0
fi

[ "$INTERACTIVE" = 1 ] && set -- -i "$@"

# shellcheck disable=SC2086
if [ -t 2 ] && [ -z "${CHARSIU_NO_PROGRESS:-}" ]; then
	# stderr through the bar, stdout straight to the terminal. The exit
	# status has to come from charsiu_run and not from awk, hence the
	# explicit plumbing rather than a plain pipeline.
	{ env $E "$RUN" "$MODEL" -c "$CTX" -t "$THREADS" "$@" 2>&1 1>&3 \
		| stage_filter >&2; } 3>&1
	exit $?
fi
exec env $E "$RUN" "$MODEL" -c "$CTX" -t "$THREADS" "$@"
