#!/usr/bin/env bash
# Strata: Flash-Next IQ3_S on qr0n's P100 build of Strata, served by tersler-strata on LLAMA_PORT.
#   strata-test.sh prep     build the engine, Python venv, weight pack and MTP head (once; no GPU needed)
#   strata-test.sh start    start tersler-strata (stops tersler-llama), wait until loaded
#   strata-test.sh stop     start tersler-llama again (stops tersler-strata)
#   strata-test.sh status   loaded?, memory, GPUs
# The GPUs serve either llama.cpp or Strata, never both. SERVE_ENGINE picks the one started at boot.
# Config: install/strata/flash-next-p100.json (qr0n's peer config with IQ3_S, --ple-io direct, 262k context, CPU image encoder).
set -euo pipefail
# shellcheck source=lib.sh
source "$(dirname "$(readlink -f "$0")")/lib.sh"
jr_load_config

SRC=$STRATA_SRC                             # git clone https://github.com/qr0n/p100-strata
SRC_COMMIT=8b1b0f5
WORK=$STRATA_WORK                           # pack, MTP head, log
CFG=$WORK/flash-next-p100.json
SHARD1=$JR_MODEL_DIR/flash-next/Qwen3.8-Flash-Next-GSQ-RCO-IQ3_S-00001-of-00002.gguf
# image encoder: ISTA's mmproj (the same file for every Flash-Next quant) + strata-vision (llama.cpp mtmd, CPU)
MMPROJ=$JR_MODEL_DIR/flash-next/mmproj-Qwen3.8-Flash-Next-BF16.gguf
MMPROJ_URL=https://huggingface.co/ISTA-DASLab/Qwen3.8-Flash-Next-GSQ-RCO-GGUF/resolve/main/mmproj-Qwen3.8-Flash-Next-BF16.gguf
MMPROJ_SHA256=b1a82259702816a5330d7bd7607cd9676b11780e79ff7348c21103ff3ce49bd0
VISION_LLAMA_COMMIT=3cf03257f219afbe7334045ff7c6a06ac68c627d   # LLAMA_CPP_COMMIT in strata/setup.py (v0.2.2)
PORT=$LLAMA_PORT
repo=$(dirname "$(dirname "$(readlink -f "$0")")")
py=$SRC/.venv/bin/python
export STRATA_GGUF_PY=$JR_PREFIX/src/llama.cpp/gguf-py   # mtp_pack/mtp_rt read GGUFs with llama.cpp's gguf-py

prep() {
    [[ -d $SRC/.git ]] || git clone https://github.com/qr0n/p100-strata "$SRC"
    [[ $(git -C "$SRC" rev-parse --short=7 HEAD) == "$SRC_COMMIT" ]] || jr_die "$SRC is not at $SRC_COMMIT"
    if ! grep -q '"+&r"(r)' "$SRC/strata/include/strata/kernels/dp4a.hpp"; then
        git -C "$SRC" apply "$repo/install/patches/strata-sm60-dp4a-alias.patch"
    fi
    if [[ ! -x $SRC/strata/build-p100/strata ]]; then
        jr_log "building the engine (sm_60, ~6-10 min)"
        (cd "$SRC" && PATH=$CUDA_HOME/bin:$PATH CUDA_ARCH=60-real JOBS=16 ./build.sh)
    fi
    if [[ ! -x $SRC/.venv/bin/pip ]]; then
        python3 -c 'import ensurepip' 2>/dev/null || sudo apt-get install -y -q python3-venv
        rm -rf "$SRC/.venv"
        python3 -m venv "$SRC/.venv"
        "$py" -m pip install -q -r "$SRC/strata/requirements.txt"
    fi
    [[ -d $WORK ]] || sudo install -d -o "$(id -un)" "$WORK"
    if [[ ! -f $WORK/packs/flash-next-iq3s/native_experts.txt ]]; then
        jr_log "weight pack (dense tensors + index; the experts stay in the GGUF)"
        (cd "$SRC/strata" && "$py" tools/iq_pack.py --gguf "$SHARD1" --out "$WORK/packs/flash-next-iq3s")
    fi
    if [[ ! -d $WORK/mtp/rt ]]; then
        jr_log "MTP head: ~5.2 GB of mtp.* tensors from Qwen/Qwen3.8-Flash-Next (HTTP range reads, SHA256-checked)"
        (cd "$SRC/strata" && "$py" tools/mtp_fetch.py fetch --out "$WORK/mtp/bf16" &&
            "$py" tools/mtp_fetch.py verify --out "$WORK/mtp/bf16" &&
            "$py" tools/mtp_pack.py --src "$WORK/mtp/bf16" --experts q2_0 --out "$WORK/mtp/mtp-q2_0.gguf" &&
            "$py" tools/mtp_rt.py --gguf "$WORK/mtp/mtp-q2_0.gguf" --out "$WORK/mtp/rt")
    fi
    if [[ ! -f $MMPROJ ]]; then
        jr_log "image encoder weights (0.9 GB)"
        curl -fL -C - --retry 5 -o "$MMPROJ.part" "$MMPROJ_URL"
        echo "$MMPROJ_SHA256  $MMPROJ.part" | sha256sum -c - >/dev/null || jr_die "checksum mismatch: $MMPROJ.part"
        mv "$MMPROJ.part" "$MMPROJ"
    fi
    if [[ ! -x $SRC/strata/build-vision/bin/strata-vision ]]; then
        local llama=$SRC/strata/third_party/llama.cpp zip=$SRC/strata/third_party/llama.cpp-${VISION_LLAMA_COMMIT:0:7}.zip
        if [[ ! -d $llama ]]; then
            curl -fL --retry 5 -o "$zip" "https://github.com/ggml-org/llama.cpp/archive/$VISION_LLAMA_COMMIT.zip"
            "$py" -c 'import sys, zipfile; zipfile.ZipFile(sys.argv[1]).extractall(sys.argv[2])' "$zip" "$SRC/strata/third_party"
            mv "$SRC/strata/third_party/llama.cpp-$VISION_LLAMA_COMMIT" "$llama"
        fi
        jr_log "building strata-vision (CPU image encoder)"
        (cd "$SRC/strata" && cmake -S tools/vision -B build-vision -G Ninja -DCMAKE_BUILD_TYPE=Release \
            -DLLAMA_DIR="$llama" -DSTRATA_PORTABLE=OFF && cmake --build build-vision --target strata-vision -j 16)
    fi
    install -m 0644 "$repo/install/strata/flash-next-p100.json" "$CFG"
    jr_log "prep done: $CFG"
}

loaded() { curl -sf "localhost:$PORT/health" | grep -q '"loaded": *true'; }

start() {
    [[ -f $CFG && -x $SRC/strata/build-p100/strata ]] || jr_die "run: $0 prep"
    jr_log "starting tersler-strata (stops tersler-llama: both need both GPUs)"
    sudo systemctl stop tersler-llama
    # drop the 27B's file pages from the cache; the expert arena needs the RAM
    sync && echo 1 | sudo tee /proc/sys/vm/drop_caches >/dev/null
    sudo systemctl start tersler-strata
    jr_log "loading (journalctl -u tersler-strata -f; engine: $WORK/strata.log)"
    for _ in $(seq 360); do loaded && { jr_log "ready on port $PORT"; status; return; }; sleep 10; done
    jr_die "not loaded after 60 min; see journalctl -u tersler-strata"
}

stop() {
    tmux kill-session -t strata 2>/dev/null || true      # first test runs used tmux
    sudo systemctl stop tersler-strata
    sudo systemctl start tersler-llama
    jr_log "Strata stopped, tersler-llama started"
}

status() {
    curl -s "localhost:$PORT/health" || echo "Strata not reachable on $PORT"
    echo
    free -g | head -2
    jr_nvsmi --query-gpu=index,memory.used,temperature.gpu,power.draw --format=csv,noheader
}

case ${1:-} in
    prep|start|stop|status) "$1" ;;
    *) sed -n '2,7p' "$0"; exit 2 ;;
esac