#!/usr/bin/env bash
# Threads-batch sweep for frank-large on dankclaw
# Benches prefill throughput at 65536 prompt tokens across threads-batch values.
set -uo pipefail

MODEL="/home/dank/models/kingjones-Qwen3.8-Flash-Next-ROCmFP4-STRIX/Qwen3.8-Flash-Next-Q4_0-ROCmFP4-STRIX.gguf"
SERVE="/home/dank/GitPackages/ROCmFPX/build-qwen38/bin/llama-server"
BENCH="/home/dank/GitPackages/ROCmFPX/build-qwen38/bin/llama-bench"
TB_LIST=(16 32 48 64 80 96 128)
ENV_FILE="/home/dank/GitPackages/ROCmFPX/build-qwen38/bin"

echo "========================================"
echo "THREADS-BATCH SWEEP — frank-large"
echo "Model: $(basename "$MODEL") ($(du -h "$MODEL" | cut -f1))"
echo "Threads-batch values: ${TB_LIST[*]}"
echo "Prompt: 65536 tokens, n-gen: 0 (prefill-only)"
echo "========================================"
echo

RESULTS=""

for tb in "${TB_LIST[@]}"; do
    echo "--- threads-batch=$tb ---"
    
    # Kill any lingering processes
    pkill -9 llama-server 2>/dev/null || true
    pkill -9 llama-bench 2>/dev/null || true
    sleep 3

    # Start server
    echo "  Starting server with --threads-batch $tb ..."
    timeout 300 env HSA_OVERRIDE_GFX_VERSION=11.5.1 HIP_VISIBLE_DEVICES=0 \
        "$SERVE" \
        --model "$MODEL" --port 7770 --jinja \
        --device ROCm0 --gpu-layers 999 --fit off \
        --threads 16 --threads-batch "$tb" \
        --parallel 2 --ctx-size 524288 \
        --batch-size 4096 --ubatch-size 4096 \
        --flash-attn on \
        --cache-type-k f16 --cache-type-v f16 \
        --cache-prompt --cache-idle-slots --metrics \
        &>/tmp/server_tb_${tb}.log &
    SERVER_PID=$!

    # Wait for ready
    READY=false
    for i in $(seq 1 300); do
        if curl -sf http://127.0.0.1:7770/health >/dev/null 2>&1; then
            READY=true
            break
        fi
        if (( i % 30 == 0 )); then
            echo "  Waiting for ready... (${i}s)"
        fi
        sleep 1
    done

    if ! $READY; then
        echo "  Server failed to start. Tail of log:"
        tail -5 /tmp/server_tb_${tb}.log
        echo "  N/A"
        RESULTS="${RESULTS}${tb},N/A,N/A,?,0\n"
        echo
        continue
    fi
    echo "  Server ready."

    # Run bench (prefill-only: n-gen=0)
    echo "  Running bench..."
    t0=$(date +%s%N)
    
    BENCH_OUTPUT=$(env HSA_OVERRIDE_GFX_VERSION=11.5.1 HIP_VISIBLE_DEVICES=0 \
        "$BENCH" \
        --model "$MODEL" --threads 16 \
        --n-prompt 65536 --n-gen 0 \
        --repetitions 3 --output csv \
        --flash-attn on --n-gpu-layers 999 2>/tmp/bench_stderr_${tb}.log)
    bexit=$?
    t1=$(date +%s%N)
    elapsed=$(( (t1 - t0) / 1000000000 ))

    # Parse CSV: find prefill row (n_prompt > 0, n_gen == 0)
    PREFILL="N/A"
    DECODE="N/A"
    CONTEXT="?"
    if [ $bexit -eq 0 ] && [ -n "$BENCH_OUTPUT" ]; then
        # Extract prefill row: n_prompt > 0 and n_gen == 0
        PREFILL=$(echo "$BENCH_OUTPUT" | python3 -c "
import csv, sys, io
reader = csv.reader(io.StringIO(sys.stdin.read()))
header = next(reader, None)
for row in reader:
    if len(row) >= 41:
        try:
            np_ = int(row[33])
            ng = int(row[34])
            avg_ts = float(row[39])
            if np_ > 0 and ng == 0:
                print(f'{avg_ts:.1f}')
                break
        except (ValueError, IndexError):
            pass
" 2>/dev/null)
        CONTEXT=$(echo "$BENCH_OUTPUT" | python3 -c "
import csv, sys, io
reader = csv.reader(io.StringIO(sys.stdin.read()))
header = next(reader, None)
for row in reader:
    if len(row) >= 41:
        try:
            np_ = int(row[33])
            ng = int(row[34])
            if np_ > 0 and ng == 0:
                print(row[40].strip() if row[40].strip() else '?')
                break
        except (ValueError, IndexError):
            pass
" 2>/dev/null)
        
        if [ -z "$PREFILL" ]; then
            PREFILL="N/A"
            CONTEXT="parse_err"
        fi
    else
        echo "  Bench failed (rc=$bexit). STDERR:"
        head -5 /tmp/bench_stderr_${tb}.log
        echo "  N/A"
    fi

    echo "  prefill=${PREFILL} tok/s  decode=${DECODE}  ctx=${CONTEXT}  time=${elapsed}s"
    RESULTS="${RESULTS}${tb},${PREFILL},${DECODE},${CONTEXT},${elapsed}\n"
    echo

    # Shutdown server
    kill $SERVER_PID 2>/dev/null || true
    wait $SERVER_PID 2>/dev/null || true
    sleep 2
done

echo "========================================"
echo "SUMMARY"
echo "========================================"
echo "tb       prefill_tok/s    decode_tok/s    ctx_util     time_s"
echo "--------------------------------------------------------"
echo -e "$RESULTS" | grep -v "^$" | while IFS=, read -r tb pt dt cu el; do
    printf "%-6s %-16s %-16s %-12s %-8s\n" "$tb" "$pt" "$dt" "$cu" "${el}s"
done

# Find best
BEST=$(echo -e "$RESULTS" | grep -v "^$" | grep -v ",N/A," | sort -t, -k2 -rn | head -1)
BEST_TB=$(echo "$BEST" | cut -d, -f1)
BEST_PT=$(echo "$BEST" | cut -d, -f2)
echo ""
echo "Best threads-batch: $BEST_TB (${BEST_PT} tok/s prefill)"

# Write results
echo -e "$RESULTS" | grep -v "^$" > /tmp/bench_tb_results.txt
echo ""
echo "Results written to /tmp/bench_tb_results.txt"

# Cleanup
pkill -9 llama-server 2>/dev/null || true
pkill -9 llama-bench 2>/dev/null || true