"""Benchmark Qwen2.5-VL-3B (Ollama, local) on the HF indian_food_images test split.
Constrained top-1: model picks one of 20 class names. Scores exact match + latency.
Writes results to bench_hf_results.jsonl and prints a summary table.
"""
import base64, io, json, time, urllib.request, os, sys

import requests
from PIL import Image

PER_CLASS = int(sys.argv[1]) if len(sys.argv) > 1 else 5
MODEL = sys.argv[2] if len(sys.argv) > 2 else "qwen2.5vl:3b"
OLLAMA = "http://localhost:11434/api/chat"  # native endpoint honors keep_alive
OUTDIR = os.path.dirname(os.path.abspath(__file__))
_tag = MODEL.replace(":", "_").replace("/", "_")
RESULTS = os.path.join(OUTDIR, f"bench_hf_results_{_tag}.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"]

PROMPT = ("You are shown a photo of a single Indian dish. Identify which ONE dish it is.\n"
    "Respond with ONLY the exact name from this list, nothing else:\n"
    + ", ".join(CLASSES) + "\nAnswer:")


IMGDIR = os.path.join(OUTDIR, "bench_hf_images")


def collect_and_download():
    """Page the test split, collect PER_CLASS urls per class, download to disk NOW
    (signed URLs expire ~15 min, so grab the bytes before classifying).
    If all expected local files already exist, reuse them and skip download."""
    os.makedirs(IMGDIR, exist_ok=True)
    expected = {c: [os.path.join(IMGDIR, f"{c}_{i}.jpg") for i in range(PER_CLASS)]
                for c in CLASSES}
    if all(os.path.exists(p) for ps in expected.values() for p in ps):
        print(f"Reusing {PER_CLASS*len(CLASSES)} cached images on disk.", flush=True)
        return expected
    buckets = {c: [] for c in CLASSES}
    offset = 0
    while offset < 941 and any(len(v) < PER_CLASS for v in buckets.values()):
        url = (f"https://datasets-server.huggingface.co/rows?dataset=rajistics/"
               f"indian_food_images&config=default&split=test&offset={offset}&length=100")
        d = json.load(urllib.request.urlopen(url))
        for r in d["rows"]:
            c = CLASSES[r["row"]["label"]]
            if len(buckets[c]) < PER_CLASS:
                buckets[c].append(r["row"]["image"]["src"])
        offset += 100
    # download all now, downscaled, to local jpgs
    paths = {c: [] for c in CLASSES}
    n = 0
    for c in CLASSES:
        for i, src in enumerate(buckets[c]):
            raw = urllib.request.urlopen(src).read()
            img = Image.open(io.BytesIO(raw)).convert("RGB")
            img.thumbnail((768, 768))
            p = os.path.join(IMGDIR, f"{c}_{i}.jpg")
            img.save(p, format="JPEG", quality=85)
            paths[c].append(p)
            n += 1
    print(f"Downloaded {n} images to disk.", flush=True)
    return paths


def to_b64(path):
    with open(path, "rb") as f:
        return base64.b64encode(f.read()).decode()


def classify(b64):
    payload = {"model": MODEL,
        "messages": [{"role": "user", "content": PROMPT, "images": [b64]}],
        "stream": False, "keep_alive": -1, "options": {"temperature": 0}}
    t0 = time.time()
    r = requests.post(OLLAMA, json=payload, timeout=300)
    dt = time.time() - t0
    raw = r.json()["message"]["content"].strip()
    pred = raw.lower().replace(" ", "_").replace("-", "_").strip(" .\n\"'")
    # keep only a known class if the model added words
    for c in CLASSES:
        if c in pred:
            pred = c
            break
    return pred, dt, raw


def warmup(first_path):
    """Pin the model in memory so the first real call isn't a cold load."""
    print("Warming up model...", flush=True)
    classify(to_b64(first_path))


def main():
    global PATHS
    print(f"Collecting + downloading {PER_CLASS} images/class from HF test split...", flush=True)
    PATHS = collect_and_download()
    total = sum(len(v) for v in PATHS.values())
    first = next(p for v in PATHS.values() for p in v)
    warmup(first)
    print(f"Got {total} images across {len(PATHS)} classes. Running...\n", flush=True)

    open(RESULTS, "w").close()
    correct = 0
    done = 0
    times = []
    misses = []
    for truth in CLASSES:
        for path in PATHS[truth]:
            try:
                b64 = to_b64(path)
                pred, dt, raw = classify(b64)
            except Exception as e:
                pred, dt, raw = f"ERROR:{e}", 0, ""
            ok = (pred == truth)
            correct += ok
            done += 1
            if dt: times.append(dt)
            if not ok: misses.append((truth, pred))
            with open(RESULTS, "a") as f:
                f.write(json.dumps({"truth": truth, "pred": pred, "ok": ok,
                                    "sec": round(dt, 1), "raw": raw}) + "\n")
            mark = "OK " if ok else "XX "
            print(f"{mark}{done:>3}/{total}  {truth:<15} -> {pred:<15} {dt:4.1f}s  "
                  f"[acc {correct}/{done} = {100*correct/done:.0f}%]", flush=True)

    avg = sum(times)/len(times) if times else 0
    print(f"\n=== RESULT ===", flush=True)
    print(f"Top-1 accuracy: {correct}/{done} = {100*correct/done:.1f}%", flush=True)
    print(f"Avg latency (warm): {avg:.1f}s", flush=True)
    if misses:
        print("Misses (truth -> predicted):", flush=True)
        for t, p in misses:
            print(f"  {t} -> {p}", flush=True)
    print("BENCH_DONE", flush=True)


if __name__ == "__main__":
    main()