import json, sys, time, threading, urllib.request
sys.path.insert(0, "/root/aiplaystf2/bot/tools")
import bench_llm as B

B.SALT = str(int(time.time()))

def one(i, out):
    t0 = time.time()
    body = {
        "model": "gemma4:e4b", "messages": B.build_messages(i), "stream": False,
        "think": False,
        "options": {"temperature": 0.85, "num_predict": 48, "num_ctx": 1024, "stop": ["\n"]},
        "keep_alive": "24h",
    }
    req = urllib.request.Request("http://localhost:11434/api/chat",
                                 data=json.dumps(body).encode(),
                                 headers={"Content-Type": "application/json"})
    r = json.loads(urllib.request.urlopen(req, timeout=300).read())
    out.append((time.time() - t0, r.get("eval_count", 0)))

for conc in (1, 2, 3):
    B.SALT = str(int(time.time())) + f"c{conc}"
    results = []
    t0 = time.time()
    batch = 6
    idx = 0
    while idx < batch:
        threads = []
        for k in range(min(conc, batch - idx)):
            t = threading.Thread(target=one, args=(idx, results)); idx += 1
            t.start(); threads.append(t)
        for t in threads: t.join()
    wall = time.time() - t0
    lat = sorted(x[0] for x in results)
    print(f"concurrency {conc}: {batch} msgs in {wall:5.1f}s "
          f"=> {batch/wall*60:5.1f} msg/min | per-msg latency mean {sum(lat)/len(lat):.1f}s p95 {lat[-1]:.1f}s")