#!/usr/bin/env python3
"""
Quality suite for the local fleet — objective auto-grading, then quality-vs-speed.

  code    : HumanEval-style problems, generated code EXECUTED against unit tests
  math    : GSM8K-style word problems, exact integer match
  mc      : multiple-choice knowledge/reasoning, letter match
  ifollow : instruction-following with programmatic checks (bullets / JSON / length)
  ocr     : multi-image text extraction (vision models only)

Grades on the model's answer (content); thinking tokens are separate in Ollama.
Also records real tok/s (aggregate eval_count / eval_duration across all calls).

Run after models are downloaded and Ollama is up:
    ./.venv/bin/python quality_bench.py
Writes quality_report.md and quality_results.json.
"""
import json, re, subprocess, tempfile, os, datetime
import benchmark as b   # reuse ollama_chat, make_ocr_image, img_user_native, list_models

# ---------------------------------------------------------------- task banks
MATH = [
    ("Natalia sold clips to 48 friends in April, then half as many in May. "
     "How many clips altogether? End with 'Answer: <number>'.", 72),
    ("Weng earns $12 an hour babysitting. Yesterday she did 50 minutes. "
     "How much did she earn? End with 'Answer: <number>'.", 10),
    ("Betty needs $100 for a wallet and has half. Her parents give $15, her "
     "grandparents twice that. How much more does she need? End with 'Answer: <number>'.", 5),
    ("James writes a 3-page letter to 2 friends twice a week. How many pages "
     "per year? End with 'Answer: <number>'.", 624),
    ("A robe takes 2 bolts of blue fiber and half that of white. How many bolts "
     "total? End with 'Answer: <number>'.", 3),
]

MC = [
    ("Which planet is largest in our solar system?\nA) Earth B) Jupiter C) Mars D) Venus\n"
     "Answer with just the letter.", "B"),
    ("What is 15% of 200?\nA) 15 B) 20 C) 30 D) 45\nAnswer with just the letter.", "C"),
    ("Antonym of 'ephemeral'?\nA) lasting B) brief C) hollow D) transient\n"
     "Answer with just the letter.", "A"),
    ("If all Bloops are Razzies and all Razzies are Lazzies, are all Bloops "
     "definitely Lazzies?\nA) yes B) no C) cannot tell D) only some\n"
     "Answer with just the letter.", "A"),
    ("Which number is prime?\nA) 21 B) 27 C) 29 D) 33\nAnswer with just the letter.", "C"),
    ("A train goes 60 km in 45 minutes. Speed in km/h?\nA) 45 B) 60 C) 75 D) 80\n"
     "Answer with just the letter.", "D"),
]

CODE = [
    ("Write a Python function `has_close_elements(numbers, threshold)` that returns True "
     "if any two numbers in the list are closer than `threshold`, else False. "
     "Return only the code in a ```python block.",
     "assert has_close_elements([1.0,2.0,3.0],0.5)==False\n"
     "assert has_close_elements([1.0,2.8,3.0,4.0,5.0,2.0],0.3)==True\n"
     "assert has_close_elements([1.0,2.0,5.9,4.0,5.0],0.95)==True\n"
     "assert has_close_elements([1.0,2.0,5.9,4.0,5.0],0.8)==False"),
    ("Write a Python function `is_palindrome(s)` that returns True if s is a palindrome "
     "ignoring case, spaces, and punctuation. Return only the code in a ```python block.",
     "assert is_palindrome('A man, a plan, a canal: Panama')==True\n"
     "assert is_palindrome('race a car')==False\n"
     "assert is_palindrome('')==True\n"
     "assert is_palindrome('Was it a car or a cat I saw?')==True"),
    ("Write a Python function `fizzbuzz(n)` returning a list of strings for 1..n: "
     "multiples of 3 -> 'Fizz', of 5 -> 'Buzz', both -> 'FizzBuzz', else the number as a "
     "string. Return only the code in a ```python block.",
     "assert fizzbuzz(5)==['1','2','Fizz','4','Buzz']\n"
     "assert fizzbuzz(15)[-1]=='FizzBuzz'\n"
     "assert fizzbuzz(3)==['1','2','Fizz']"),
]

IFOLLOW = [
    ("List exactly three primary colors. Format your ENTIRE response as exactly three "
     "bullet points, each line starting with '- ' and nothing else.", "bullets"),
    ("Output ONLY a valid JSON object with keys \"name\" (string) and \"age\" (number) "
     "for a person named Alice who is 30. No markdown, no extra text.", "json"),
    ("Reply with a single sentence of fewer than 10 words describing the sun.", "short"),
]

OCR_TEXTS = ["INVOICE-7X42", "Total: $1,284.50", "SKU-99Z-KAPPA"]
VISION = ("qwen3-vl", "gemma4", "gemma3")

# ---------------------------------------------------------------- extractors / graders
def extract_int(text):
    m = re.findall(r"[Aa]nswer[^0-9\-]*(-?\d[\d,]*)", text)
    nums = re.findall(r"-?\d[\d,]*", text)
    cand = m[-1] if m else (nums[-1] if nums else None)
    if cand is None:
        return None
    try:
        return int(float(cand.replace(",", "")))
    except Exception:
        return None

def extract_letter(text):
    m = re.search(r"(?:answer|correct)\D{0,12}([A-D])\b", text, re.I)
    if m:
        return m.group(1).upper()
    allm = re.findall(r"\b([A-D])\b", text)
    return allm[-1] if allm else None

def extract_code(content, thinking):
    for src in (content, content + "\n" + thinking):
        m = re.search(r"```(?:python)?\s*(.*?)```", src, re.S)
        if m and "def " in m.group(1):
            return m.group(1)
        i = src.find("def ")
        if i >= 0:
            return src[i:]
    return ""

def run_code(code, tests):
    prog = code + "\n\n" + tests + "\nprint('__OK__')\n"
    with tempfile.NamedTemporaryFile("w", suffix=".py", delete=False) as f:
        f.write(prog); path = f.name
    try:
        r = subprocess.run(["python3", path], capture_output=True, text=True, timeout=12)
        return 1 if r.returncode == 0 and "__OK__" in r.stdout else 0
    except Exception:
        return 0
    finally:
        os.unlink(path)

def check_if(kind, text):
    t = text.strip()
    if kind == "bullets":
        lines = [l for l in t.splitlines() if l.strip()]
        return 1 if len(lines) == 3 and all(l.strip().startswith("- ") for l in lines) else 0
    if kind == "json":
        m = re.search(r"\{.*\}", re.sub(r"```(json)?", "", t), re.S)
        if not m:
            return 0
        try:
            d = json.loads(m.group(0))
            return 1 if str(d.get("name", "")).lower() == "alice" and int(d.get("age")) == 30 else 0
        except Exception:
            return 0
    if kind == "short":
        n = len(t.split())
        return 1 if 0 < n < 10 else 0
    return 0

# ---------------------------------------------------------------- per-model run
class Meter:
    def __init__(self):
        self.ec = 0.0; self.ed = 0.0
    def call(self, model, messages, num_predict):
        # think disabled for a fast, direct-answer, apples-to-apples baseline;
        # fall back if a model rejects the param
        try:
            d = b.ollama_chat(model, messages, num_predict, think=False)
        except Exception:
            d = b.ollama_chat(model, messages, num_predict)
        self.ec += d.get("eval_count", 0) or 0
        self.ed += (d.get("eval_duration", 0) or 0) / 1e9
        msg = d.get("message") or {}
        return msg.get("content", "") or "", msg.get("thinking", "") or ""
    def tps(self):
        return self.ec / self.ed if self.ed else 0.0


def run_model(model, ocr_imgs):
    vis = any(v in model for v in VISION)
    m = Meter()
    scores = {}

    got = sum(1 for q, a in MATH
              if extract_int(m.call(model, [b.user(q)], 320)[0]) == a)
    scores["math"] = got / len(MATH)

    got = sum(1 for q, a in MC
              if extract_letter(m.call(model, [b.user(q)], 48)[0]) == a)
    scores["mc"] = got / len(MC)

    got = 0
    for q, tests in CODE:
        c, th = m.call(model, [b.user(q)], 512)
        got += run_code(extract_code(c, th), tests)
    scores["code"] = got / len(CODE)

    got = sum(check_if(kind, m.call(model, [b.user(q)], 160)[0]) for q, kind in IFOLLOW)
    scores["ifollow"] = got / len(IFOLLOW)

    if vis and ocr_imgs:
        got = 0
        for text, img in zip(OCR_TEXTS, ocr_imgs):
            c, th = m.call(model, [b.img_user_native("What text is shown in this image?", img)], 256)
            blob = (c + " " + th).replace("-", "").replace(" ", "").lower()
            got += 1 if text.replace("-", "").replace(" ", "").replace(",", "").replace("$", "").lower()[:8] \
                in blob.replace(",", "").replace("$", "") else 0
        scores["ocr"] = got / len(OCR_TEXTS)

    text_cats = [scores[k] for k in ("math", "mc", "code", "ifollow")]
    scores["overall"] = sum(text_cats) / len(text_cats)
    scores["tps"] = m.tps()
    return scores


def main():
    models = b.list_models()
    ocr_imgs = None
    try:
        ocr_imgs = [b.make_ocr_image(t) for t in OCR_TEXTS]
    except Exception as e:
        print("pillow unavailable, OCR skipped:", e)

    results = {}
    for i, model in enumerate(models, 1):
        print(f"[{i}/{len(models)}] {model} ...", flush=True)
        try:
            results[model] = run_model(model, ocr_imgs)
            s = results[model]
            print(f"    overall {s['overall']*100:4.0f}%  "
                  f"(code {s['code']*100:.0f} math {s['math']*100:.0f} "
                  f"mc {s['mc']*100:.0f} if {s['ifollow']*100:.0f}"
                  f"{' ocr '+str(int(s.get('ocr',0)*100)) if 'ocr' in s else ''})  "
                  f"{s['tps']:.1f} tok/s", flush=True)
        except Exception as e:
            print(f"    ERROR: {e}", flush=True)
            results[model] = {"error": str(e)}

    with open("quality_results.json", "w") as f:
        json.dump(results, f, indent=2)

    ts = datetime.datetime.now().strftime("%Y-%m-%d %H:%M")
    ranked = sorted([(m, s) for m, s in results.items() if "overall" in s],
                    key=lambda x: x[1]["overall"], reverse=True)
    with open("quality_report.md", "w") as f:
        f.write(f"# Quality benchmark — {ts}\n\n")
        f.write("Objective auto-grading. code=executed unit tests, math=GSM8K-style exact, "
                "mc=multiple-choice, ifollow=programmatic checks, ocr=multi-image text.\n\n")
        f.write("| Model | Overall | Code | Math | MC | Instr | OCR | tok/s |\n")
        f.write("|---|--:|--:|--:|--:|--:|--:|--:|\n")
        for m, s in ranked:
            ocr = f"{s['ocr']*100:.0f}" if "ocr" in s else "—"
            f.write(f"| `{m}` | **{s['overall']*100:.0f}%** | {s['code']*100:.0f} | "
                    f"{s['math']*100:.0f} | {s['mc']*100:.0f} | {s['ifollow']*100:.0f} | "
                    f"{ocr} | {s['tps']:.1f} |\n")
    print("\nWrote quality_report.md and quality_results.json")


if __name__ == "__main__":
    main()