"""Benchmark moondream on the same 200 HF images.
moondream is a captioner, not a constrained classifier, so the fair metric is:
ask "what food is this?" and check whether the true dish name (or a common
synonym) appears in the caption. Documents the speed/accuracy tradeoff vs Qwen2.5-VL.
"""
import base64, json, time, glob, os, sys, requests

MODEL = "moondream"
URL = "http://localhost:11434/api/chat"
OUTDIR = os.path.dirname(os.path.abspath(__file__))
IMGDIR = os.path.join(OUTDIR, "bench_hf_images")
RESULTS = os.path.join(OUTDIR, "bench_hf_results_moondream.jsonl")
PROMPT = "Name the food in this image."  # terse prompts make moondream emit garbage; captions work

# true-class -> substrings that count as a correct identification (lenient, fair to a captioner)
SYN = {
    "burger": ["burger"],
    "butter_naan": ["naan"],
    "chai": ["chai", "tea"],
    "chapati": ["chapati", "chapatti", "roti", "flatbread"],
    "chole_bhature": ["chole", "chana", "bhatur", "chickpea"],
    "dal_makhani": ["dal", "daal", "lentil"],
    "dhokla": ["dhokla"],
    "fried_rice": ["fried rice", "rice"],
    "idli": ["idli"],
    "jalebi": ["jalebi"],
    "kaathi_rolls": ["kathi", "kaathi", "roll", "wrap"],
    "kadai_paneer": ["paneer", "cottage cheese"],
    "kulfi": ["kulfi", "ice cream"],
    "masala_dosa": ["dosa"],
    "momos": ["momo", "dumpling"],
    "paani_puri": ["pani puri", "paani", "golgappa", "puri"],
    "pakode": ["pakora", "pakoda", "pakode", "fritter"],
    "pav_bhaji": ["pav", "bhaji"],
    "pizza": ["pizza"],
    "samosa": ["samosa"],
}
CLASSES = list(SYN.keys())


def classify(b64):
    p = {"model": MODEL, "messages": [{"role": "user", "content": PROMPT, "images": [b64]}],
         "stream": False, "keep_alive": -1, "options": {"temperature": 0}}
    t = time.time()
    r = requests.post(URL, json=p, timeout=180)
    dt = time.time() - t
    return r.json()["message"]["content"].strip(), dt


def scored(truth, caption):
    c = caption.lower()
    return any(s in c for s in SYN[truth])


def main():
    imgs = {ct: sorted(glob.glob(os.path.join(IMGDIR, f"{ct}_*.jpg"))) for ct in CLASSES}
    total = sum(len(v) for v in imgs.values())
    # warmup
    first = next(p for v in imgs.values() for p in v)
    print("warming up moondream...", flush=True)
    classify(base64.b64encode(open(first, "rb").read()).decode())
    print(f"Running {total} images...\n", flush=True)

    # resume: if results already exist, skip the pairs already done and append
    prior = []
    if os.path.exists(RESULTS):
        prior = [json.loads(l) for l in open(RESULTS, encoding="utf-8")]
    skip = len(prior)
    correct = sum(r["ok"] for r in prior)
    done = skip
    times = [r["sec"] for r in prior if r.get("sec")]
    if skip:
        print(f"Resuming: {skip} already done, {correct}/{skip} correct so far.\n", flush=True)

    fmode = "a" if skip else "w"
    if fmode == "w":
        open(RESULTS, "w").close()
    seen = 0
    for truth in CLASSES:
        for path in imgs[truth]:
            if seen < skip:
                seen += 1
                continue
            b64 = base64.b64encode(open(path, "rb").read()).decode()
            try:
                cap, dt = classify(b64)
            except Exception as e:
                cap, dt = f"ERROR:{e}", 0
            ok = scored(truth, cap)
            correct += ok; done += 1
            if dt: times.append(dt)
            with open(RESULTS, "a", encoding="utf-8") as f:
                f.write(json.dumps({"truth": truth, "ok": ok, "sec": round(dt, 1),
                                    "caption": cap}, ensure_ascii=False) + "\n")
            mark = "OK " if ok else "XX "
            safe = cap[:38].encode("ascii", "replace").decode()
            print(f"{mark}{done:>3}/{total}  {truth:<15} {safe:<40} {dt:4.1f}s "
                  f"[{correct}/{done}={100*correct/done:.0f}%]", flush=True)

    avg = sum(times)/len(times) if times else 0
    print(f"\n=== RESULT (moondream, keyword-match metric) ===", flush=True)
    print(f"Dish identified: {correct}/{done} = {100*correct/done:.1f}%", flush=True)
    print(f"Avg latency (warm): {avg:.1f}s", flush=True)
    print("BENCH_DONE", flush=True)


if __name__ == "__main__":
    main()