#!/usr/bin/env bash
# Decode throughput benchmark for frank-default-large.
# Measures generation tok/s with 8K+ token prompts and 256 generated tokens.
set -euo pipefail

SERVE=/home/dank/GitPackages/ROCmFPX/build-qwen38/bin/llama-server
MODEL=/home/dank/models/kingjones-Qwen3.8-Flash-Next-ROCmFP4-STRIX/Qwen3.8-Flash-Next-Q4_0-ROCmFP4-STRIX.gguf
PROMPT_FILE=/tmp/bench-prompts/prompt_code.txt
PORT=7770
N_GEN=256
SERVER_DIR=/home/dank/GitPackages/ROCmFPX/build-qwen38/bin
ROCM_PATH=/opt/rocm
LIB_PATH="$SERVER_DIR:$ROCM_PATH/lib:$ROCM_PATH/lib/rocm_sysdeps/lib:$ROCM_PATH/lib/llvm/lib"

kill_all() {
    pkill -9 -f llama-server 2>/dev/null || true
    sleep 3
}

wait_ready() {
    for i in $(seq 1 600); do
        if curl -s http://127.0.0.1:$PORT/health 2>/dev/null | grep -q ok; then
            return 0
        fi
        sleep 1
    done
    return 1
}

run_bench() {
    local tb=$1
    local logfile=/tmp/server_tb${tb}.log
    
    kill_all
    
    # Start server
    env HSA_OVERRIDE_GFX_VERSION=11.5.1 HIP_VISIBLE_DEVICES=0 \
        LD_LIBRARY_PATH="$LIB_PATH" ROCM_PATH="$ROCM_PATH" \
        nohup $SERVE \
        --model "$MODEL" --port $PORT --jinja \
        --device ROCm0 --gpu-layers 999 --fit off \
        --threads 16 --threads-batch $tb \
        --parallel 2 --ctx-size 131072 \
        --batch-size 4096 --ubatch-size 4096 \
        --flash-attn on \
        --cache-type-k f16 --cache-type-v f16 \
        --cache-prompt --cache-idle-slots --metrics \
        --temp 0.7 --top-p 0.8 --top-k 20 --min-p 0.0 \
        --reasoning on --reasoning-format deepseek \
        --reasoning-budget -1 > "$logfile" 2>&1 &
    
    if ! wait_ready; then
        echo "TB=$tb: FAILED (server not ready)"
        return 1
    fi
    
    # Warmup (small request)
    curl -s -m 60 http://127.0.0.1:$PORT/v1/chat/completions \
        -H "Content-Type: application/json" \
        -d '{"model":"frank-default-large","messages":[{"role":"user","content":"Hi"}],"max_tokens":16,"stream":false}' > /dev/null 2>&1 || true
    
    # Build 8K prompt
    local prompt
    prompt=$(python3 -c "
import json
with open('$PROMPT_FILE') as f: text=f.read()
words=text.split()[:8200]
print(json.dumps(' '.join(words)))
")
    
    echo "TB=$tb: Starting benchmark..."
    
    # Benchmark
    local t0 t1 resp pt gt wall gtps
    t0=$(date +%s%N)
    resp=$(curl -s -m 600 http://127.0.0.1:$PORT/v1/chat/completions \
        -H "Content-Type: application/json" \
        -d "{\"model\":\"frank-default-large\",\"messages\":[{\"role\":\"user\",\"content\":$prompt}],\"max_tokens\":$N_GEN,\"stream\":false}" 2>&1)
    t1=$(date +%s%N)
    
    if ! echo "$resp" | python3 -c "import sys,json; json.load(sys.stdin)" 2>/dev/null; then
        echo "TB=$tb: FAILED (HTTP error: $(echo "$resp" | head -1))"
        return 1
    fi
    
    pt=$(echo "$resp" | python3 -c "import sys,json; d=json.load(sys.stdin); print(d['usage']['prompt_tokens'])")
    gt=$(echo "$resp" | python3 -c "import sys,json; d=json.load(sys.stdin); print(d['usage']['completion_tokens'])")
    wall=$(python3 -c "print(round(($t1-$t0)/1000000000,1))")
    gtps=$(python3 -c "print(round($N_GEN/($t1-$t0)*1000000000,2))")
    
    echo "TB=$tb: pt=$pt gt=$gt wall=${wall}s gen=${gtps} tok/s"
    
    # Extract prefill timing from server log
    local prefill_time
    prefill_time=$(grep -oP 'prompt eval time = \K[0-9.]+' "$logfile" | head -1 || echo "N/A")
    if [ "$prefill_time" != "N/A" ]; then
        echo "TB=$tb: prefill=$(python3 -c "print(round($N_GEN/($wall),2))") tok/s"
    fi
    
    kill_all
    sleep 2
    
    echo "$tb,$gt,$wall,$gtps,$pt" >> /tmp/bench_decode_results.csv
}

# Main
echo "=== Decode throughput benchmark ==="
echo "Prompt: ~8K tokens, Generation: $N_GEN tokens"
echo ""

# Clear results
echo "tb,gen_tokens,wall_s,gen_tps,prompt_tokens" > /tmp/bench_decode_results.csv

# Run benchmarks
for tb in 16 32; do
    run_bench $tb
    echo ""
done

# Summary
echo "=== SUMMARY ==="
echo "Prefill throughput (from earlier benchmark):"
echo "  TB=16: 373.5 tok/s"
echo "  TB=32: 354.4 tok/s"
echo ""
echo "Decode throughput (generation tok/s with 8K prompt):"
tail -n +2 /tmp/bench_decode_results.csv | while IFS=, read tb gt wall gtps pt; do
    echo "  TB=$tb: $gt tokens in ${wall}s = ${gtps} tok/s (pt=$pt)"
done