#!/usr/bin/env python3
"""Task-class x concurrency decode benchmark for the AEON serve (OpenAI chat API, vLLM metrics).
gen tok/s = completion tokens / (wall - TTFT) per stream; aggregate = sum tokens / max wall (decode window).
tok/pass = delta(spec accepted)/delta(drafts) + 1 from /metrics."""
import json, os, sys, time, threading, urllib.request, statistics, random
BASE=os.environ.get("BASE","http://localhost:8000"); MODEL=os.environ.get("MODEL","aeon"); OUT=int(os.environ.get("OUT","512"))
THINK=os.environ.get("THINK","0")=="1"
TOPICS=["a lighthouse keeper in 1890s Norway","a courier robot in a flooded Bangkok","a violinist who loses her hearing","a mountain village's last winter","two brothers repairing a boat","a night-shift nurse in Lagos","a chess prodigy in exile","a beekeeper during a drought"]
TASKS={
 "prose": lambda i: f"Write a vivid literary short story (at least 700 words) about {TOPICS[i%8]}. No headings, flowing prose only.",
 "code":  lambda i: f"Write a complete, well-documented Python module implementing an LRU cache with TTL expiry, thread safety, statistics, and a small CLI demo. Variant #{i}: name the class Cache{i}. Code only, no prose.",
 "list":  lambda i: f"List 120 distinct {['programming concepts','world cities','chemical compounds','historical events','musical instruments','mathematical theorems','plant species','cooking techniques'][i%8]}, numbered, each with a one-line description.",
 "essay": lambda i: f"Write a structured argumentative essay (about 800 words, with an introduction, three sections with headers, and a conclusion) on: '{['Should cities ban private cars?','Is remote work good for innovation?','Should AI systems have a right to refuse tasks?','Is nuclear power essential for decarbonization?','Should voting be mandatory?','Do social networks make us lonelier?','Should space exploration be privatized?','Is a four-day work week viable?'][i%8]}'",
}
def metrics():
    t=urllib.request.urlopen(BASE+"/metrics",timeout=10).read().decode(); d={}
    for l in t.splitlines():
        for k in ("spec_decode_num_drafts_total","spec_decode_num_accepted_tokens_total"):
            if l.startswith("vllm:"+k+"{"): d[k]=float(l.rsplit(" ",1)[1])
    return d
def one(prompt, res, idx):
    body={"model":MODEL,"messages":[{"role":"user","content":prompt}],"max_tokens":OUT,"stream":True,
          "stream_options":{"include_usage":True},"chat_template_kwargs":{"enable_thinking":THINK}}
    body.update(json.loads(os.environ.get("SAMPLING","{}")))
    req=urllib.request.Request(BASE+"/v1/chat/completions",data=json.dumps(body).encode(),headers={"Content-Type":"application/json"})
    t0=time.time(); ttft=None; ntok=0
    with urllib.request.urlopen(req,timeout=1800) as r:
        for line in r:
            line=line.decode().strip()
            if not line.startswith("data:") or line=="data: [DONE]": continue
            ev=json.loads(line[5:])
            if ev.get("usage"): ntok=ev["usage"]["completion_tokens"]
            ch=ev.get("choices")
            d=(ch[0].get("delta") or {}) if ch else {}
            if ttft is None and (d.get("content") or d.get("reasoning_content") or d.get("reasoning")): ttft=time.time()-t0
    wall=time.time()-t0; res[idx]=(ntok,ttft or wall,wall)
def run(task,conc):
    m0=metrics(); res=[None]*conc; th=[threading.Thread(target=one,args=(TASKS[task](i+conc*7),res,i)) for i in range(conc)]
    T=time.time(); [t.start() for t in th]; [t.join() for t in th]; W=time.time()-T; m1=metrics()
    toks=sum(r[0] for r in res); per=[r[0]/(r[2]-r[1]) for r in res if r[2]>r[1]] or [float("nan")]
    g=lambda m,k: m.get(k,0.0)
    dr=g(m1,"spec_decode_num_drafts_total")-g(m0,"spec_decode_num_drafts_total"); ac=g(m1,"spec_decode_num_accepted_tokens_total")-g(m0,"spec_decode_num_accepted_tokens_total")
    tpp=(ac/dr+1) if dr else float("nan"); ttft=statistics.median(r[1] for r in res)
    return dict(task=task,conc=conc,agg=toks/W,per=statistics.mean(per),ttft=ttft,tpp=tpp,tok=toks/conc)
if __name__=="__main__":
    tasks=sys.argv[1].split(",") if len(sys.argv)>1 else list(TASKS); concs=[int(c) for c in (sys.argv[2] if len(sys.argv)>2 else "1,2,4,8").split(",")]
    run("prose",1)  # warm
    print(f"thinking={'on' if THINK else 'off'} out={OUT} sampling=server default (temp 0.6/top_p .95/top_k 20/rep 1.05)")
    print(f"{'task':6} {'C':>2} {'agg t/s':>8} {'per-stream':>10} {'TTFT':>6} {'tok/pass':>8} {'tok/req':>7}")
    for task in tasks:
        for c in concs:
            r=run(task,c); print(f"{r['task']:6} {r['conc']:>2} {r['agg']:8.1f} {r['per']:10.1f} {r['ttft']:6.2f} {r['tpp']:8.2f} {r['tok']:7.0f}",flush=True)