"""Parallel benchmark harness: fires N concurrent requests at llama-server's
parallel slots (default 4) instead of one-at-a-time. Same disambiguated v2 prompt
and scoring; cuts wall-time ~Nx. Per-image latency is still recorded (though
concurrent requests share the GPU, so per-call time rises while throughput wins).

Usage: python bench_parallel.py [workers]   (default 4)
"""
import base64, json, time, glob, os, sys, requests
from concurrent.futures import ThreadPoolExecutor, as_completed

URL = "http://127.0.0.1:8080/v1/chat/completions"
WORKERS = int(sys.argv[1]) if len(sys.argv) > 1 else 4
OUTDIR = os.path.dirname(os.path.abspath(__file__))
IMGDIR = os.path.join(OUTDIR, "bench_hf_images")
RESULTS = os.path.join(OUTDIR, "bench_hf_results_parallel.jsonl")

CLASSES = ["burger","butter_naan","chai","chapati","chole_bhature","dal_makhani",
    "dhokla","fried_rice","idli","jalebi","kaathi_rolls","kadai_paneer","kulfi",
    "masala_dosa","momos","paani_puri","pakode","pav_bhaji","pizza","samosa"]
DESC = {
    "burger": "bun with a patty",
    "butter_naan": "teardrop-shaped leavened flatbread, glossy with butter",
    "chai": "milky tea in a cup or glass",
    "chapati": "plain thin unleavened round flatbread, no filling",
    "chole_bhature": "chickpea curry served WITH a large puffy fried bread (bhatura)",
    "dal_makhani": "creamy dark black-lentil curry (only if clearly lentils)",
    "dhokla": "steamed yellow spongy savoury cake, cut in squares",
    "fried_rice": "stir-fried rice with visible separate grains",
    "idli": "plain white round steamed rice cakes",
    "jalebi": "bright orange crispy spiral-shaped sweet",
    "kaathi_rolls": "a rolled paratha/wrap around a filling, NOT a curry",
    "kadai_paneer": "white paneer cubes in a thick tomato-pepper gravy",
    "kulfi": "dense frozen milk dessert, often on a stick",
    "masala_dosa": "large thin folded crispy crepe with potato filling",
    "momos": "pleated steamed or fried dumplings",
    "paani_puri": "small round hollow crispy puris (golgappa)",
    "pakode": "irregular deep-fried fritters/clumps",
    "pav_bhaji": "mashed red-orange vegetable curry served WITH soft bread rolls (only if mashed veg + pav)",
    "pizza": "flat bread base with cheese and toppings",
    "samosa": "triangular fried pastry with filling",
}
PROMPT = ("You are shown a photo of a single Indian dish. Identify which ONE dish it is.\n"
    "Choose the best match from this list (name: description):\n"
    + "\n".join(f"- {c}: {DESC[c]}" for c in CLASSES)
    + "\n\nDo not default to pav_bhaji or dal_makhani unless the photo clearly matches "
      "their descriptions. Respond with ONLY the exact class name (left of the colon), nothing else.\nAnswer:")


def classify(path, truth):
    with open(path, "rb") as f:
        b64 = base64.b64encode(f.read()).decode()
    body = {"messages": [{"role": "user", "content": [
        {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{b64}"}},
        {"type": "text", "text": PROMPT}]}],
        "temperature": 0, "max_tokens": 30}
    t = time.time()
    try:
        r = requests.post(URL, json=body, timeout=300)
        raw = r.json()["choices"][0]["message"]["content"].strip()
    except Exception as e:
        raw = f"ERROR:{e}"
    dt = time.time() - t
    pred = raw.lower().replace(" ", "_").replace("-", "_").strip(" .\n\"'")
    for c in CLASSES:
        if c in pred:
            pred = c
            break
    return {"truth": truth, "pred": pred, "ok": pred == truth, "sec": round(dt, 1), "raw": raw}


def main():
    jobs = []
    for c in CLASSES:
        for p in sorted(glob.glob(os.path.join(IMGDIR, f"{c}_*.jpg"))):
            jobs.append((p, c))
    print(f"{len(jobs)} images, {WORKERS} concurrent workers...", flush=True)

    open(RESULTS, "w").close()
    correct = done = 0
    wall0 = time.time()
    with ThreadPoolExecutor(max_workers=WORKERS) as ex:
        futs = [ex.submit(classify, p, t) for p, t in jobs]
        for fut in as_completed(futs):
            res = fut.result()
            correct += res["ok"]; done += 1
            with open(RESULTS, "a", encoding="utf-8") as f:
                f.write(json.dumps(res) + "\n")
            if done % 20 == 0 or done == len(jobs):
                print(f"  {done}/{len(jobs)}  acc {correct}/{done}={100*correct/done:.0f}%", flush=True)
    wall = time.time() - wall0

    rows = [json.loads(l) for l in open(RESULTS, encoding="utf-8")]
    t = sorted(r["sec"] for r in rows if r["sec"] > 0)
    print(f"\n=== RESULT (parallel x{WORKERS}) ===", flush=True)
    print(f"Top-1 accuracy: {correct}/{done} = {100*correct/done:.1f}%", flush=True)
    print(f"Wall time: {wall:.0f}s for {done} images ({wall/done:.2f}s/image throughput)", flush=True)
    print(f"Per-call latency (concurrent): median={t[len(t)//2]:.1f} mean={sum(t)/len(t):.1f} max={t[-1]:.1f}", flush=True)
    print("BENCH_DONE", flush=True)


if __name__ == "__main__":
    main()