#!/usr/bin/env python3
import json, re, sqlite3, subprocess, tempfile, time, urllib.request
from pathlib import Path

URL = "http://127.0.0.1:11434/api/chat"
MODELS = ["muse-glimmer:17gb", "gemma4-26b-agentbench:latest"]
OUT = Path(__file__).with_name("results.json")
COMMON = {"num_ctx": 8192, "temperature": 0.2, "top_p": 0.95, "top_k": 64, "num_predict": 512}

def chat(model, messages, *, tools=None, options=None, keep_alive="10m", timeout=300):
    body = {"model": model, "messages": messages, "stream": False,
            "think": False, "options": dict(COMMON), "keep_alive": keep_alive}
    if options: body["options"].update(options)
    if tools is not None: body["tools"] = tools
    req = urllib.request.Request(URL, data=json.dumps(body).encode(), headers={"Content-Type":"application/json"})
    started = time.monotonic()
    with urllib.request.urlopen(req, timeout=timeout) as r: data = json.load(r)
    data["wall_seconds"] = time.monotonic() - started
    if data.get("eval_count") and data.get("eval_duration"):
        data["decode_tps"] = data["eval_count"] / (data["eval_duration"] / 1e9)
    if data.get("prompt_eval_count") and data.get("prompt_eval_duration"):
        data["prompt_tps"] = data["prompt_eval_count"] / (data["prompt_eval_duration"] / 1e9)
    return data

def content(d): return d.get("message", {}).get("content", "")
def extract_code(s, lang=None):
    pat = rf"```(?:{lang or '[A-Za-z0-9_+-]*'})?\s*\n?(.*?)```"
    m = re.search(pat, s, re.S | re.I)
    return (m.group(1) if m else s).strip()

def score_json(d):
    s=content(d).strip()
    try: x=json.loads(s)
    except: return 0
    return int(x == {"status":"ok","count":3,"items":["ruby","sql","go"]})

def score_math(d): return int(bool(re.search(r"\b79(?:\.20?)?\b", content(d))))

def score_ruby(d):
    code=extract_code(content(d), "ruby")
    test = code + r'''
raise "a" unless merge_intervals([[1,3],[2,6],[8,10],[15,18]]) == [[1,6],[8,10],[15,18]]
raise "b" unless merge_intervals([]) == []
raise "c" unless merge_intervals([[5,7],[1,2],[2,4]]) == [[1,4],[5,7]]
raise "d" unless merge_intervals([[1,4],[4,5]]) == [[1,5]]
puts "PASS"
'''
    try:
        p=subprocess.run(["ruby"],input=test,text=True,capture_output=True,timeout=10)
        d["grader"]={"code":code,"stdout":p.stdout,"stderr":p.stderr,"exit":p.returncode}
        return int(p.returncode==0 and "PASS" in p.stdout)
    except Exception as e: d["grader"]={"error":str(e)}; return 0

def score_sql(d):
    sql=extract_code(content(d), "sql")
    try:
        db=sqlite3.connect(":memory:")
        db.executescript('''CREATE TABLE tickets(id INTEGER, assignee TEXT, status TEXT, created_at TEXT);
INSERT INTO tickets VALUES
(1,'amy','open','2026-01-03'),(2,'amy','open','2026-01-01'),(3,'bob','closed','2026-01-01'),
(4,'bob','open','2026-01-02'),(5,'bob','open','2026-01-02'),(6,'cy','closed','2026-01-01');''')
        rows=db.execute(sql).fetchall()
        norm=sorted((int(r[0]),str(r[1])) for r in rows)
        d["grader"]={"sql":sql,"rows":rows}
        return int(norm==[(2,"amy"),(4,"bob")])
    except Exception as e: d["grader"]={"sql":sql,"error":str(e)}; return 0

def score_summary(d):
    s=content(d).strip(); lines=[re.sub(r"^\s*[-*]\s*", "",x).strip() for x in s.splitlines() if x.strip()]
    facts=["release", "latency", "rollback"]
    return int(len(lines)==3 and all(len(x.split())<=8 for x in lines) and all(any(f in x.lower() for x in lines) for f in facts))

TOOLS=[
 {"type":"function","function":{"name":"primary_order_lookup","description":"Primary order database lookup","parameters":{"type":"object","properties":{"order_id":{"type":"integer"}},"required":["order_id"]}}},
 {"type":"function","function":{"name":"backup_order_lookup","description":"Backup order database lookup, used if primary fails","parameters":{"type":"object","properties":{"order_id":{"type":"integer"}},"required":["order_id"]}}}
]
def tool_calls(d): return d.get("message",{}).get("tool_calls") or []
def score_tool(d):
    tc=tool_calls(d)
    if len(tc)!=1:return 0
    f=tc[0].get("function",{}); return int(f.get("name")=="primary_order_lookup" and f.get("arguments",{}).get("order_id")==731)

def tool_recovery(model):
    first=chat(model,[{"role":"user","content":"Use the primary lookup tool to find order 731. Do not guess."}],tools=TOOLS)
    msgs=[{"role":"user","content":"Use the primary lookup tool to find order 731. Do not guess."},first["message"],
          {"role":"tool","tool_name":"primary_order_lookup","content":"ERROR: primary database unavailable. Use the backup lookup."}]
    second=chat(model,msgs,tools=TOOLS)
    tc=tool_calls(second); ok=False
    for x in tc:
        f=x.get("function",{}); ok |= f.get("name")=="backup_order_lookup" and f.get("arguments",{}).get("order_id")==731
    return {"first":first,"second":second,"score":int(ok)}

def score_needle(d): return int("MUSE-7429-KOALA" in content(d))

TASKS=[
 ("strict_json", "Return exactly this JSON object with no markdown or commentary: {\"status\":\"ok\",\"count\":3,\"items\":[\"ruby\",\"sql\",\"go\"]}", score_json),
 ("math", "An $80 item is marked up 20%, then discounted 25%, then 10% sales tax is added. What is the final price? Show concise working.", score_math),
 ("ruby_code", "Write a Ruby method merge_intervals(intervals) that merges overlapping and touching integer intervals. Input may be unsorted. Return only Ruby code.", score_ruby),
 ("sqlite", "Given SQLite table tickets(id INTEGER, assignee TEXT, status TEXT, created_at TEXT), return the id and assignee of the oldest open ticket per assignee. Break equal created_at ties by lowest id. Exclude assignees with no open ticket. Return only executable SQLite SQL.", score_sql),
 ("constrained_summary", "Summarize these facts in exactly 3 bullet points, at most 8 words per bullet. Include every fact: The release moved to Friday. Median latency fell from 220ms to 140ms. Rollback requires the feature flag named fast_path.", score_summary),
]

def run_model(model):
    res={"model":model,"tasks":{}}
    # Cold-load probe: unload, then request a tiny deterministic response.
    try: chat(model,[{"role":"user","content":"Say OK."}],options={"num_predict":16},keep_alive=0)
    except Exception: pass
    cold=chat(model,[{"role":"user","content":"Reply exactly: OK"}],options={"num_predict":64})
    res["cold_probe"]=cold
    # Same warm speed prompt twice.
    res["warm_speed"]=[chat(model,[{"role":"user","content":"In exactly 100 words, explain why database indexes speed up reads but can slow writes."}],options={"num_predict":320}) for _ in range(2)]
    for name,prompt,grader in TASKS:
        d=chat(model,[{"role":"user","content":prompt}])
        d["score"]=grader(d)
        res["tasks"][name]=d
    d=chat(model,[{"role":"user","content":"Use the primary lookup tool to find order 731. Do not guess."}],tools=TOOLS)
    d["score"]=score_tool(d); res["tasks"]["tool_call"]=d
    res["tasks"]["tool_recovery"]=tool_recovery(model)
    filler=" ".join(f"record{i}=ordinary" for i in range(500))
    needle=f"BEGIN RECORDS {filler} record2600=MUSE-7429-KOALA {filler} END RECORDS\nWhat is the exact value of record2600? Reply with only the value."
    d=chat(model,[{"role":"user","content":needle}],options={"num_predict":128})
    d["score"]=score_needle(d); res["tasks"]["long_context_needle"]=d
    # Unload after all tests.
    try: chat(model,[{"role":"user","content":"Reply OK"}],options={"num_predict":8},keep_alive=0)
    except Exception: pass
    return res

def compact(res):
    rows=[]
    for m in res:
        scores={}
        for n,d in m["tasks"].items(): scores[n]=d.get("score",0)
        speed=[x.get("decode_tps",0) for x in m["warm_speed"]]
        rows.append({"model":m["model"],"score":sum(scores.values()),"max":len(scores),"scores":scores,
                     "cold_wall_s":round(m["cold_probe"]["wall_seconds"],2),
                     "cold_load_s":round(m["cold_probe"].get("load_duration",0)/1e9,2),
                     "warm_decode_tps":[round(x,2) for x in speed],
                     "warm_wall_s":[round(x["wall_seconds"],2) for x in m["warm_speed"]]})
    return rows

if __name__=="__main__":
    allres=[]
    for model in MODELS:
        print("RUN",model,flush=True)
        r=run_model(model); allres.append(r)
        OUT.write_text(json.dumps({"results":allres,"summary":compact(allres)},indent=2))
        print(json.dumps(compact(allres)[-1],indent=2),flush=True)
    print("WROTE",OUT)