#!/usr/bin/env bash
# Bench the Kairic Edge "fast greedy (default)" profile.
# Usage: ./bench-fast-greedy.sh
#
# Config: reasoning off, argmax fastpath=1, deterministic sampler.
# Sends a coding prompt, captures full stream, computes decode tokens and time.
set -euo pipefail

HOST="127.0.0.1"
PORT=7770
ALIAS="frank-default-large"

# A coding prompt representative of the card's 164-task suite.
PROMPT='### Instruction
Write a Python class that implements a priority queue using a binary heap. Include methods for enqueue, dequeue, peek, and is_empty. Add type hints, docstrings, and a usage example. The heap should support arbitrary comparable elements.
### Response'

DECODE_TOKENS=0
START_SEC=0
END_SEC=0

echo "=== Fast-greedy benchmark (FASTGREEDY=1) ==="
echo "Model: $ALIAS on $HOST:$PORT"
echo "Prompt tokens: $(echo "$PROMPT" | wc -c) chars"
echo ""

# Use /v1/completions with streaming, capture tokens+time from usage.
RESPONSE=$(curl -sf http://$HOST:$PORT/v1/completions \
  -H 'Content-Type: application/json' \
  -d "{
    \"model\": \"$ALIAS\",
    \"prompt\": $(python3 -c "import json,sys; print(json.dumps(sys.stdin.read()))" <<< "$PROMPT"),
    \"max_tokens\": 512,
    \"temperature\": 0,
    \"top_p\": 1.0,
    \"top_k\": 0,
    \"min_p\": 0.0,
    \"n\": 1,
    \"stream\": true
  }" 2>/dev/null)

# Extract final usage from the last JSON line (finish_reason=stop or length).
# Also count output text tokens.
LAST_LINE=$(echo "$RESPONSE" | grep 'finish_reason' | tail -1)

# Parse timing from the last chunk (has full usage)
if echo "$LAST_LINE" | grep -q 'timings'; then
    COMP_TOKENS=$(echo "$LAST_LINE" | python3 -c "
import sys, json
data = json.loads(sys.stdin.read())
print(data.get('usage', {}).get('completion_tokens', 0))
print(data.get('timings', {}).get('predicted_per_second', 0))
print(data.get('timings', {}).get('prompt_per_second', 0))
" 2>/dev/null)
    echo "$COMP_TOKENS" | {
        read CT
        read TPS
        read PPS
        echo "Completion tokens: $CT"
        echo "Predicted (decode) tokens/sec: $TPS"
        echo "Prompt tokens/sec: $PPS"
    }
else
    echo "WARNING: No timing data found in response"
    echo "$LAST_LINE" | head -5
fi

# Count the actual text output
OUTPUT_TEXT=$(echo "$RESPONSE" | grep '"text"' | head -1 | sed 's/.*"text":"//;s/","index.*//')
OUTPUT_TOKENS=$(echo -n "$OUTPUT_TEXT" | wc -w)
echo "Output text (rough token count): ~$OUTPUT_TOKENS words"
echo ""

# Check journal for draft acceptance and final timing
echo "=== Draft acceptance & final timing ==="
sleep 2
ssh -p 2222 dank@dankclaw "journalctl -u frank-default-large --since '1 minute ago' --no-pager 2>/dev/null | grep -E 'draft acceptance|print_timing.*n_decoded|statistics.*draft-mtp' | tail -4" 2>/dev/null || echo "(could not read journal)"