#!/usr/bin/env python3
"""Measure generation throughput on workloads where ngram speculative decoding
can actually pay off: output that largely echoes the prompt (refactor / rewrite /
RAG-style restatement), plus a free-form control that it cannot help with."""
import json, argparse, urllib.request

CODE = '''
class RingBuffer:
    def __init__(self, capacity):
        self.capacity = capacity
        self.data = [None] * capacity
        self.head = 0
        self.tail = 0
        self.size = 0

    def push(self, item):
        if self.size == self.capacity:
            raise BufferError("ring buffer is full")
        self.data[self.tail] = item
        self.tail = (self.tail + 1) % self.capacity
        self.size += 1

    def pop(self):
        if self.size == 0:
            raise IndexError("pop from empty ring buffer")
        item = self.data[self.head]
        self.data[self.head] = None
        self.head = (self.head + 1) % self.capacity
        self.size -= 1
        return item
'''

WORKLOADS = {
    # Best case: output is mostly a copy of the input with small edits.
    "rewrite": "Repeat this Python class verbatim, changing only the class name "
               "to CircularBuffer. Output only code, no commentary:\n" + CODE,
    # Realistic coding task: structurally predictable, not a copy.
    "extend":  "Here is a Python class:\n" + CODE +
               "\nRewrite it with an added peek() method and type hints. Output only code.",
    # Control: novel prose, nothing for an ngram matcher to hit.
    "prose":   "Write a short original story about a lighthouse keeper who collects storms.",
}


def run(port, name, prompt, n_predict, no_think):
    body = {"messages": [{"role": "user", "content": prompt}],
            "max_tokens": n_predict, "temperature": 0}
    if no_think:
        body["chat_template_kwargs"] = {"enable_thinking": False}
    req = urllib.request.Request(f"http://localhost:{port}/v1/chat/completions",
                                 data=json.dumps(body).encode(),
                                 headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=1800) as r:
        d = json.load(r)
    t = d.get("timings", {})
    tps = t.get("predicted_per_second") or 0
    n = t.get("predicted_n") or 0
    print(f"  {name:<10} {n:>5} tok @ {tps:>6.2f} t/s", flush=True)
    return tps


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--port", type=int, default=18010)
    ap.add_argument("--label", default="run")
    ap.add_argument("--n-predict", type=int, default=400)
    ap.add_argument("--no-think", action="store_true",
                    help="disable reasoning so the measurement is the answer, not the thinking")
    a = ap.parse_args()

    print(f"[{a.label}]", flush=True)
    for name, prompt in WORKLOADS.items():
        try:
            run(a.port, name, prompt, a.n_predict, a.no_think)
        except Exception as e:
            print(f"  {name:<10} FAILED {type(e).__name__}: {e}", flush=True)