#!/bin/bash
#
# Copyright (C) 2026 Danilo M. <danix@danix.xyz>
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 2 as
# published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# The one runnable check for ai-state.sh. It puts stub curl, pgrep,
# llamachat and imggen on PATH and a stub assistant.sh beside them, then runs
# the real script, so nothing here touches the live stack. T_* variables pick
# the state each stub reports.
#
# Usage: ./test-ai-state.sh   (exit 0 = all passed)

set -u

here="$(cd "$(dirname "$0")" && pwd)"
stub="$(mktemp -d)"
trap 'rm -rf "$stub"' EXIT

# The last argument is the URL. /slots refuses a probe without
# autoload=false, because the real server would load the model for it.
cat > "$stub/curl" <<'STUB'
#!/bin/bash
url="${!#}"
case "$url" in
    */models)
        case "${T_LLAMA:-down}" in
            down) exit 7 ;;
            unloaded) echo '{"data":[{"id":"A","status":{"value":"unloaded"}}]}' ;;
            *) echo '{"data":[{"id":"A","status":{"value":"unloaded"}},{"id":"Gemma","status":{"value":"loaded"}}]}' ;;
        esac ;;
    */slots)
        [[ " $* " == *" autoload=false "* ]] || exit 99
        [[ " $* " == *" model=Gemma "* ]] || exit 22
        if [[ "$T_LLAMA" == busy ]]; then
            echo '[{"id":0,"is_processing":false},{"id":1,"is_processing":true}]'
        else
            echo '[{"id":0,"is_processing":false}]'
        fi ;;
    *) exit 6 ;;
esac
STUB

cat > "$stub/pgrep" <<'STUB'
#!/bin/bash
case "$*" in
    *sd-server*)
        [[ -n "${T_SD:-}" ]] || exit 1
        echo "4242 sd-server --listen-port 7860 --diffusion-model /data/SD/z_image_turbo-Q8_0.gguf --vae /data/SD/vae/flux1-ae.safetensors" ;;
    *fanfictioner*)
        [[ -n "${T_FANFIC:-}" ]] || exit 1
        echo 31337 ;;
    *) exit 1 ;;
esac
STUB

cat > "$stub/llamachat" <<'STUB'
#!/bin/bash
[[ "$1" == --ping && -n "${T_CHAT:-}" ]]
STUB

# The real imggen exits 0 either way and says which on stdout.
cat > "$stub/imggen" <<'STUB'
#!/bin/bash
[[ "$1" == status ]] || exit 1
if [[ -n "${T_IMG:-}" ]]; then echo '{"model": "realvis", "ready": true}'; else echo "daemon down"; fi
STUB

cat > "$stub/assistant.sh" <<'STUB'
#!/bin/bash
[[ "$1" == status && -n "${T_ASSIST:-}" ]]
STUB

chmod +x "$stub"/*
log="$stub/sd-server.log"
: > "$log"

pass=0
fail=0

check() {
    local name="$1" want="$2" got="$3"
    if [[ "$want" == "$got" ]]; then
        pass=$((pass + 1))
    else
        fail=$((fail + 1))
        printf 'FAIL: %s\n  want: %q\n  got:  %q\n' "$name" "$want" "$got"
    fi
}

# One line of the protocol, tab joined.
l() { local IFS=$'\t'; printf '%s\n' "$*"; }

run() {
    env PATH="$stub:$PATH" AI_ASSISTANT="$stub/assistant.sh" AI_SD_LOG="$log" "$@" \
        bash "$here/ai-state.sh"
}

apps_down="$(l app assistant stopped -; l app chat stopped -; l app imggen stopped -; l app fanfic stopped -)"

want="$(l engine llama down idle -; l engine sd down idle -)
$apps_down"
check "everything down" "$want" "$(run)"
run >/dev/null
check "everything down exits zero" 0 $?

got="$(run T_LLAMA=unloaded | head -n1)"
check "llama up, no model loaded" "$(l engine llama up idle -)" "$got"

got="$(run T_LLAMA=idle | head -n1)"
check "llama loaded, idle" "$(l engine llama up idle Gemma)" "$got"

got="$(run T_LLAMA=busy | head -n1)"
check "llama loaded, generating" "$(l engine llama up busy Gemma)" "$got"

touch "$log"
got="$(run T_SD=1 | sed -n 2p)"
check "sd running, log just written" "$(l engine sd up busy z_image_turbo-Q8_0.gguf)" "$got"

touch -d '1 minute ago' "$log"
got="$(run T_SD=1 | sed -n 2p)"
check "sd running, log quiet" "$(l engine sd up idle z_image_turbo-Q8_0.gguf)" "$got"

want="$(l app assistant running -; l app chat running -; l app imggen running realvis; l app fanfic running 31337)"
got="$(run T_ASSIST=1 T_CHAT=1 T_IMG=1 T_FANFIC=1 | tail -n4)"
check "every app running" "$want" "$got"

printf '\n%d passed, %d failed\n' "$pass" "$fail"
[[ "$fail" -eq 0 ]]