#!/usr/bin/env bash
# bench-frank-dragoon.sh — 64k-context benchmark for the frank-default-large stack.
#
# Measures generation throughput of Qwen3.6-27B Q8_K_P + DFlash + ngram-mod on
# a single backend at 64k context. It starts llama-server on a scratch port,
# drives one generation through the OpenAI-compatible API with a ~64k-token
# prompt, and prints tokens/sec.
#
# Usage:
#   ./scripts/bench-frank-dragoon.sh <ROCm0|Vulkan0> [--generations N]
#
# Example:
#   ./scripts/bench-frank-dragoon.sh ROCm0
#   ./scripts/bench-frank-dragoon.sh Vulkan0 --generations 3
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
DEVICE="${1:?usage: bench-frank-dragoon.sh <ROCm0|Vulkan0> [--generations N]}"
GENERATIONS=1
if [ "${2:-}" = "--generations" ]; then
    GENERATIONS="${3:?--generations needs a count}"
fi

MODEL_DIR=/home/dank/models
MODEL_FILE=Qwen3.6-27B-Uncensored-HauhauCS-Aggressive-Q8_K_P.gguf
DRAFT_FILE=Qwen3.6-27B-DFlash-Q8_0.gguf
LLAMA_SERVER=/home/dank/llama.cpp/build/bin/llama-server
MODEL="$MODEL_DIR/$MODEL_FILE"
DRAFT="$MODEL_DIR/$DRAFT_FILE"

PORT=7890
CTX=65536
PARALLEL=1
BATCH=2048
UBATCH=512
OUT_TOKENS=256

for f in "$MODEL" "$DRAFT"; do
    if [[ ! -f "$f" ]]; then
        echo "missing required file: $f" >&2
        exit 1
    fi
done

echo "=== 64k benchmark: device=$DEVICE model=$MODEL_FILE ==="

# --- start the server ----------------------------------------------------
"$LLAMA_SERVER" \
    --model "$MODEL" \
    --spec-draft-model "$DRAFT" \
    --spec-type "draft-dflash,ngram-mod" \
    --spec-draft-n-max 8 \
    --spec-ngram-mod-n-min 48 \
    --spec-ngram-mod-n-max 64 \
    --spec-ngram-mod-n-match 24 \
    --device "$DEVICE" \
    --gpu-layers all \
    --cache-type-k q8_0 \
    --cache-type-v f16 \
    --ctx-size "$CTX" \
    --parallel "$PARALLEL" \
    --batch-size "$BATCH" \
    --ubatch-size "$UBATCH" \
    --alias bench \
    --host 127.0.0.1 \
    --port "$PORT" \
    > "/tmp/frankbench-$DEVICE.log" 2>&1 &
SERVER_PID=$!
echo "server pid $SERVER_PID on 127.0.0.1:$PORT"

trap 'kill $SERVER_PID 2>/dev/null; wait $SERVER_PID 2>/dev/null' EXIT

# --- wait for ready ------------------------------------------------------
echo -n "waiting for server"
ready=0
for i in $(seq 1 400); do
    if curl -sf "http://127.0.0.1:$PORT/health" >/dev/null 2>&1; then
        echo " [ready]"
        ready=1
        break
    fi
    echo -n "."
    sleep 2
done
if [ "$ready" != "1" ]; then
    echo ""
    echo "FAIL: server did not become ready"
    tail -25 "/tmp/frankbench-$DEVICE.log"
    exit 1
fi

# --- drive N generations and measure --------------------------------------
python3 - "$DEVICE" "$CTX" "$OUT_TOKENS" "$GENERATIONS" "$PORT" <<'PY'
import sys, time, json, urllib.request, statistics

device       = sys.argv[1]
ctx          = int(sys.argv[2])
out_tokens   = int(sys.argv[3])
generations  = int(sys.argv[4])
port         = sys.argv[5]

# ~64k-token prompt. One prose unit is ~110 chars ≈ 28 tokens for this model.
# We aim for a prompt of ~60k tokens so generation stays within the 64k ctx.
unit = ("The careful study of natural language reveals patterns of meaning "
        "that emerge from the deliberate arrangement of words and phrases. ")
# approximate: unit is 160 chars, ~40 tokens; repeat to ~1500 units = ~60k tokens
prompt = unit * 1500

def req(prompt, n):
    body = json.dumps({
        "prompt": prompt,
        "max_tokens": n,
        "temperature": 0.0,
        "stream": False,
    }).encode()
    r = urllib.request.Request(f"http://127.0.0.1:{port}/v1/completions",
                               data=body, headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(r, timeout=1800) as resp:
        return json.load(resp)

# prefill timing
t0 = time.time()
first = req(prompt, 0)
t_pre = time.time() - t0
used = first.get("usage", {})
prompt_tokens = used.get("prompt_tokens", 0)
print(f"prefill: {prompt_tokens} tokens in {t_pre:.2f}s "
      f"({prompt_tokens/t_pre:.1f} tok/s)")

# generation timing
rates = []
total_gentok = 0
total_time = 0.0
for i in range(generations):
    t0 = time.time()
    out = req(prompt, out_tokens)
    dt = time.time() - t0
    uit = out.get("usage", {})
    comp = uit.get("completion_tokens", 0)
    gt = dt - t_pre  # approximate generation time = total - prefill
    r = comp / gt if gt > 0 else 0
    rates.append(r)
    total_gentok += comp
    total_time += gt
    print(f"gen {i+1}: {comp} tokens in {gt:.2f}s ({r:.1f} tok/s)")

avg = statistics.mean(rates) if rates else 0
print(f"\nRESULT device={device} prompt_tokens={prompt_tokens} "
      f"avg_gen_tok_s={avg:.1f} over {generations} runs")
PY