#!/usr/bin/env bash
# Benchmark n-gram speculative decoding vs no-spec for frank-default-large.
#
# This tests prompt-processing throughput at 64k context with two distinct
# prompts: one code-heavy and one prose-heavy. Uses llama-server on scratch
# ports with the OpenAI-compatible API to measure prompt_eval_duration_ms.
#
# Run on Dankclaw after stopping the production service:
#   sudo systemctl stop frank-default-large
#   ./scripts/bench-ngram-mod.sh
#   sudo systemctl start frank-default-large

set -euo pipefail

SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
ENV_FILE="${FRANK_LARGE_ENV_FILE:-$SCRIPT_DIR/frank-default-large.env}"
if [[ -f "$ENV_FILE" ]]; then
    set -a; source "$ENV_FILE"; set +a
fi

: "${LLAMA_BIN:=/home/dank/GitPackages/ROCmFPX/build-qwen38/bin/llama-server}"
: "${MODEL_DIR:=/home/dank/models/kingjones-Qwen3.8-Flash-Next-ROCmFP4-STRIX}"
: "${MODEL_FILE:=Qwen3.8-Flash-Next-Q4_0-ROCmFP4-STRIX.gguf}"
: "${CHAT_TEMPLATE_FILE:=$SCRIPT_DIR/chat-template.jinja}"
: "${CHAT_TEMPLATE_KWARGS:={}}"
: "${BATCH_SIZE:=4096}"
: "${UBATCH_SIZE:=4096}"
: "${CTX_SIZE:=65536}"
: "${NUM_THREADS:=16}"
: "${THREAD_BATCH:=32}"
: "${CACHE_TYPE_K:=f16}"
: "${CACHE_TYPE_V:=f16}"
: "${FIT:=off}"
: "${TEMPERATURE:=0.7}"
: "${TOP_P:=0.8}"
: "${TOP_K:=20}"
: "${MIN_P:=0.0}"
: "${REASONING:=on}"
: "${REASONING_FORMAT:=deepseek}"

MODEL="$MODEL_DIR/$MODEL_FILE"
HOSTNAME=$(hostname 2>/dev/null || echo "dankclaw")

# Prompt files (generated by generate-prompts.py or passed via args)
PROMPT_CODE="${1:-}"
PROMPT_PROSE="${2:-}"

if [[ -z "$PROMPT_CODE" || -z "$PROMPT_PROSE" ]]; then
    echo "ERROR: Need two prompt file paths."
    echo "Usage: $0 <prompt_code.txt> <prompt_prose.txt>"
    echo "Generate with: python3 scripts/generate-prompts.py"
    exit 1
fi

if [[ ! -r "$MODEL" ]]; then
    echo "missing or unreadable model: $MODEL" >&2
    exit 1
fi

SERVER_DIR="$(cd -- "$(dirname -- "$LLAMA_BIN")" && pwd)"
ROCM_PATH="${ROCM_PATH:-/opt/rocm}"
LIBRARY_PATH="$SERVER_DIR:$ROCM_PATH/lib:$ROCM_PATH/lib/rocm_sysdeps/lib:$ROCM_PATH/lib/llvm/lib${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}"

# Scratch ports
PORT_NGRAM=7771
PORT_BASELINE=7772

cleanup() {
    echo "Cleaning up servers..."
    # Kill any leftover llama-server on our ports
    for port in $PORT_NGRAM $PORT_BASELINE; do
        fuser -k "$port"/tcp 2>/dev/null || true
    done
}
trap cleanup EXIT

wait_for_server() {
    local port=$1
    local max_wait=120
    local elapsed=0
    while (( elapsed < max_wait )); do
        if curl -s "http://127.0.0.1:$port/health" | grep -q '"running"'; then
            return 0
        fi
        sleep 1
        (( elapsed++ ))
    done
    echo "ERROR: Server on port $port did not start within ${max_wait}s" >&2
    return 1
}

get_prompt_tokens() {
    local file=$1
    # Count tokens approximately (1 token ~ 4 chars for English/code)
    wc -c < "$file" | awk '{printf "%d", $1/4}'
}

run_single_test() {
    local port=$1
    local spec_type="$2"
    local prompt_file=$3
    local label=$4
    local ngram_n_max="${5:-24}"
    local ngram_n_min="${6:-4}"
    
    echo "=== $label ==="
    echo "  spec-type: $spec_type"
    echo "  ngram-mod n-max: $ngram_n_max (if enabled)"
    echo "  ngram-mod n-min: $ngram_n_min (if enabled)"
    echo "  prompt: $prompt_file ($(wc -c < "$prompt_file") bytes)"
    
    # Build command
    local cmd=(
        /usr/bin/env
        ROCM_PATH="$ROCM_PATH"
        HIP_VISIBLE_DEVICES=0
        LD_LIBRARY_PATH="$LIBRARY_PATH"
        HSA_OVERRIDE_GFX_VERSION="${HSA_OVERRIDE_GFX_VERSION:-11.5.1}"
        "$LLAMA_BIN"
        --model "$MODEL"
        --chat-template-file "$CHAT_TEMPLATE_FILE"
        --chat-template-kwargs "$CHAT_TEMPLATE_KWARGS"
        --host 127.0.0.1
        --port "$port"
        --jinja
        --device ROCm0
        --gpu-layers 999
        --parallel 1
        --ctx-size "$CTX_SIZE"
        --batch-size "$BATCH_SIZE"
        --ubatch-size "$UBATCH_SIZE"
        --fit "$FIT"
        --flash-attn on
        --cache-type-k "$CACHE_TYPE_K"
        --cache-type-v "$CACHE_TYPE_V"
        --threads "$NUM_THREADS"
        --threads-batch "$THREAD_BATCH"
        --cache-prompt
        --cache-idle-slots
        --metrics
        --temp "$TEMPERATURE"
        --top-p "$TOP_P"
        --top-k "$TOP_K"
        --min-p "$MIN_P"
        --reasoning "$REASONING"
        --reasoning-format "$REASONING_FORMAT"
        --reasoning-budget -1
    )
    
    # Add spec flags
    if [[ "$spec_type" == "ngram-mod" ]]; then
        cmd+=(
            --spec-type "ngram-mod"
            --spec-ngram-mod-n-max "$ngram_n_max"
            --spec-ngram-mod-n-min "$ngram_n_min"
        )
    else
        cmd+=(
            --spec-type "none"
        )
    fi
    
    # Start server with stderr to log
    local log_file="/tmp/ngram-bench-${label//[^a-zA-Z0-9]/-}.log"
    echo "  Starting server (log: $log_file)..."
    
    "${cmd[@]}" 2> "$log_file" &
    local server_pid=$!
    
    if ! wait_for_server "$port"; then
        echo "ERROR: Server failed to start" >&2
        kill "$server_pid" 2>/dev/null || true
        return 1
    fi
    
    echo "  Server ready."
    
    # Send prompt via API
    echo "  Sending prompt..."
    local prompt_text
    prompt_text=$(cat "$prompt_file")
    
    local response
    response=$(curl -s -X POST "http://127.0.0.1:$port/v1/completions" \
        -H "Content-Type: application/json" \
        -d "{
            \"model\": \"frank-default-large\",
            \"prompt\": \"$prompt_text\",
            \"max_tokens\": 128,
            \"temperature\": 0,
            \"stream\": false
        }")
    
    # Extract metrics
    local prompt_tokens prompt_eval_duration_ms total_duration_ms
    prompt_tokens=$(echo "$response" | python3 -c "import sys,json; d=json.load(sys.stdin); print(d['prompt_eval_count'])" 2>/dev/null || echo "0")
    prompt_eval_duration_ms=$(echo "$response" | python3 -c "import sys,json; d=json.load(sys.stdin); print(d['prompt_eval_duration_ms'])" 2>/dev/null || echo "0")
    total_duration_ms=$(echo "$response" | python3 -c "import sys,json; d=json.load(sys.stdin); print(d['total_duration_ms'])" 2>/dev/null || echo "0")
    
    # Calculate throughput
    local prompt_tps=0
    if [[ "$prompt_eval_duration_ms" -gt 0 ]]; then
        prompt_tps=$(python3 -c "print(f'{($prompt_tokens / $prompt_eval_duration_ms) * 1000:.1f}')")
    fi
    
    echo "  Prompt tokens:  $prompt_tokens"
    echo "  Prompt eval:    ${prompt_eval_duration_ms}ms"
    echo "  Prompt TPS:     $prompt_tps tok/s"
    echo "  Total duration: ${total_duration_ms}ms"
    
    # Kill server
    kill "$server_pid" 2>/dev/null || true
    sleep 2
    
    echo "  Done."
    echo "---"
    
    # Output structured results for parsing
    echo "RESULT|$label|$prompt_file|$prompt_tokens|$prompt_eval_duration_ms|$total_duration_ms|$prompt_tps"
}

echo "╔══════════════════════════════════════════════════════════╗"
echo "║  N-gram Mod Benchmark — frank-default-large             ║"
echo "║  STRIX ROCmFP4 · 64k context · two prompts             ║"
echo "╚══════════════════════════════════════════════════════════╝"
echo ""
echo "Host: $HOSTNAME"
echo "Model: $MODEL"
echo "Build: $(basename "$LLAMA_BIN")"
echo "Context: $CTX_SIZE"
echo ""

# Run baseline (no spec)
echo "=== BASELINE (no speculative decoding) ==="
echo ""

BASELINE_CODE_RESULT=$(run_single_test $PORT_BASELINE "none" "$PROMPT_CODE" "Baseline-Code" 24 4)
BASELINE_PROSE_RESULT=$(run_single_test $PORT_NGRAM "none" "$PROMPT_PROSE" "Baseline-Prose" 24 4)

echo ""
echo "=== NGRAM-MOD (speculative decoding enabled) ==="
echo ""

# Run n-gram mod
NGRAM_CODE_RESULT=$(run_single_test $PORT_NGRAM "ngram-mod" "$PROMPT_CODE" "Ngram-Code" 24 4)
NGRAM_PROSE_RESULT=$(run_single_test $PORT_BASELINE "ngram-mod" "$PROMPT_PROSE" "Ngram-Prose" 24 4)

echo ""
echo "╔══════════════════════════════════════════════════════════╗"
echo "║  RESULTS SUMMARY                                       ║"
echo "╚══════════════════════════════════════════════════════════╝"
echo ""

# Parse and display results
python3 << 'PYEOF'
import sys

results = []
for line in sys.stdin:
    line = line.strip()
    if not line.startswith("RESULT|"):
        continue
    parts = line.split("|")
    if len(parts) != 8:
        continue
    label, prompt_file, prompt_tokens, prompt_eval_ms, total_ms, tps = parts[1], parts[2], int(parts[3]), int(parts[4]), int(parts[5]), float(parts[6])
    results.append({
        'label': label,
        'prompt_file': prompt_file,
        'prompt_tokens': prompt_tokens,
        'prompt_eval_ms': prompt_eval_ms,
        'total_ms': total_ms,
        'tps': tps
    })

# Compare baseline vs ngram for each prompt type
for code_result in results:
    if not code_result['label'].startswith('Ngram-'):
        continue
    prompt_type = code_result['label'].replace('Ngram-', '').replace('Ngram', 'Baseline').replace('-', '-')
    baseline = next((r for r in results if r['label'] == 'Baseline-' + prompt_type), None)
    if not baseline:
        continue
    
    spec = code_result['label'] == 'Ngram-' + prompt_type
    
    print(f"Prompt: {prompt_type}")
    print(f"  Tokens:   {code_result['prompt_tokens']}")
    if spec:
        print(f"  Spec:     ngram-mod (n_max=24, n_min=4)")
        print(f"  Prompt:   {code_result['prompt_eval_ms']}ms  |  {code_result['tps']:>10.1f} tok/s")
        baseline_tps = baseline['tps']
        if baseline_tps > 0:
            improvement = ((code_result['tps'] - baseline_tps) / baseline_tps) * 100
            print(f"  Speedup:  {improvement:+.1f}% vs baseline ({baseline_tps:.1f} tok/s)")
    else:
        print(f"  Spec:     none (baseline)")
        print(f"  Prompt:   {baseline['prompt_eval_ms']}ms  |  {baseline['tps']:>10.1f} tok/s")
    print()
PYEOF

echo ""
echo "Raw data lines above (RESULT|... format) can be parsed for automation."
echo ""