"""Banc 'N agents Claude Code' : N sessions en parallele, chacune un prompt
propre (~PROMPT_K jetons, sale), puis T tours qui ajoutent un tool_result de
~3 k jetons. Flux SSE : TTFT, debit de decodage par flux, cache_read. Verdict
= debit agrege en regime mixte prefill/decodage, ce que les bancs courts ne
voient pas."""
import json, sys, time, uuid, threading, requests
B="https://y26flhzhufrkb3-8080.proxy.runpod.net"
H={"User-Agent":"curl/8","x-api-key":"x","anthropic-version":"2023-06-01","content-type":"application/json","accept":"text/event-stream"}
N=int(sys.argv[1]) if len(sys.argv)>1 else 6
T=int(sys.argv[2]) if len(sys.argv)>2 else 4
PROMPT_K=int(sys.argv[3]) if len(sys.argv)>3 else 50
OUT=int(sys.argv[4]) if len(sys.argv)>4 else 300
ADD_K=int(sys.argv[5]) if len(sys.argv)>5 else 1   # k jetons ajoutes par tour (tool_result) ; 30 = "agent lecteur"
MODEL="claude-ornith[1m]"
res=[]; lock=threading.Lock()
def stream(msgs, agent, tour):
    body={"model":MODEL,"max_tokens":OUT,"stream":True,"messages":msgs,"thinking":{"type":"enabled","budget_tokens":OUT//2}}
    t0=time.time(); first=None; n=0; usage={}
    with requests.post(B+"/v1/messages",headers=H,json=body,stream=True,timeout=900) as r:
        if r.status_code!=200:
            with lock: print(f"A{agent} t{tour} HTTP {r.status_code}",flush=True)
            return None
        for line in r.iter_lines():
            if not line or not line.startswith(b"data:"): continue
            try: d=json.loads(line[5:])
            except Exception: continue
            ty=d.get("type")
            if ty=="content_block_delta":
                if first is None: first=time.time()
                n+=1
            elif ty=="message_start": usage.update(d["message"].get("usage",{}))
            elif ty=="message_delta": usage.update(d.get("usage",{}))
    t1=time.time(); ttft=(first or t1)-t0; dec=(t1-(first or t1))
    out=usage.get("output_tokens",n); cr=usage.get("cache_read_input_tokens",0) or 0; inp=usage.get("input_tokens",0)
    row=dict(agent=agent,tour=tour,total=inp+cr,cache_read=cr,miss=inp,ttft=ttft,out=out,tps=out/dec if dec>0 else 0,wall=t1-t0)
    with lock:
        res.append(row); print(f"A{agent} t{tour} ctx={inp+cr:6d} miss={inp:6d} TTFT={ttft:5.1f}s out={out:4d} decode={row['tps']:5.1f} j/s  ({t1-t0:5.1f}s)",flush=True)
    return row
def agent(i):
    salt=uuid.uuid4().hex
    doc=("Dossier %s, pièce %%d : le contrôleur valide la saisie, journalise l'appel, puis délègue au service métier qui persiste l'entité et publie un événement. " % salt)
    doc="".join(doc % k for k in range(PROMPT_K*1000//45))
    msgs=[{"role":"user","content":doc+"\nRelis ce dossier et dis-moi ce qu'il manque."}]
    for t in range(T):
        stream(msgs,i,t)
        msgs.append({"role":"assistant","content":[{"type":"text","text":"Je lis la pièce suivante."}]})
        msgs.append({"role":"user","content":("Pièce complémentaire %d-%d : "%(i,t)+"ligne de journal applicatif, horodatée, niveau INFO, corrélée à la requête. ")*(150*ADD_K)})
t0=time.time(); th=[threading.Thread(target=agent,args=(i,)) for i in range(N)]
[x.start() for x in th]; [x.join() for x in th]
wall=time.time()-t0
tot_out=sum(r["out"] for r in res); tot_miss=sum(r["miss"] for r in res); tot_ctx=sum(r["total"] for r in res)
print(f"\nBILAN N={N} T={T} prompt~{PROMPT_K}k : {len(res)} requetes en {wall:.0f}s | sortie {tot_out} jetons = {tot_out/wall:.0f} j/s agrege | prefill reel {tot_miss} jetons = {tot_miss/wall:.0f} j/s | cache {100*(1-tot_miss/max(1,tot_ctx)):.0f} % | TTFT moy {sum(r['ttft'] for r in res)/len(res):.1f}s | decode moy/flux {sum(r['tps'] for r in res)/len(res):.1f} j/s")
json.dump(res,open("agent_bench_last.json","w"))