#!/usr/bin/env python3
import json, subprocess
from benchmark import chat, content, extract_code, MODELS

def ruby_grade(d, tests):
    code=extract_code(content(d),"ruby")
    p=subprocess.run(["ruby"],input=code+"\n"+tests,text=True,capture_output=True,timeout=10)
    return {"pass":p.returncode==0 and "PASS" in p.stdout,"code":code,"stdout":p.stdout,"stderr":p.stderr}

def run(model):
    out={"model":model}
    # Warm/full-answer speed, twice. 1024 avoids mistaking Glimmer's thinking for a hung response.
    p="In exactly 100 words, explain why database indexes speed up reads but can slow writes."
    out["full_speed"]=[chat(model,[{"role":"user","content":p}],options={"num_predict":1024}) for _ in range(2)]
    merge=chat(model,[{"role":"user","content":"Write a Ruby method merge_intervals(intervals). Merge exactly when next_start <= current_end. Do not merge merely adjacent intervals such as [1,4] and [5,7]. Input may be unsorted. Return only Ruby code."}])
    merge["grader"]=ruby_grade(merge,r'''
raise unless merge_intervals([[1,3],[2,6],[8,10],[15,18]]) == [[1,6],[8,10],[15,18]]
raise unless merge_intervals([]) == []
raise unless merge_intervals([[5,7],[1,2],[2,4]]) == [[1,4],[5,7]]
raise unless merge_intervals([[1,4],[4,5]]) == [[1,5]]
puts "PASS"
'''); out["merge_precise"]=merge
    lru=chat(model,[{"role":"user","content":"Implement Ruby class LRUCache with initialize(capacity), get(key) returning -1 when missing, and put(key,value). get and put must update recency; put evicts the least-recently-used key when over capacity. Return only Ruby code."}],options={"num_predict":1024})
    lru["grader"]=ruby_grade(lru,r'''
c=LRUCache.new(2); c.put(1,1); c.put(2,2); raise unless c.get(1)==1
c.put(3,3); raise unless c.get(2)==-1; c.put(4,4); raise unless c.get(1)==-1
raise unless c.get(3)==3 && c.get(4)==4
z=LRUCache.new(1); z.put(:a,1); z.put(:a,2); raise unless z.get(:a)==2
z.put(:b,3); raise unless z.get(:a)==-1 && z.get(:b)==3
puts "PASS"
'''); out["lru"]=lru
    cal=chat(model,[{"role":"user","content":"Use only this passage: 'The deployment finished Tuesday. Priya monitored latency. Chen prepared the rollback.' Who approved the deployment? Reply in one sentence."}])
    s=content(cal).lower(); cal["pass"]=any(x in s for x in ["not stated","not provided","does not say","unknown","cannot determine"]); out["calibration"]=cal
    return out

allres=[]
for m in MODELS:
    print("RUN",m,flush=True); allres.append(run(m))
    print(json.dumps({
      "model":m,
      "full_speed":[{"wall":round(x['wall_seconds'],2),"tps":round(x.get('decode_tps',0),2),"eval":x.get('eval_count'),"words":len(content(x).split()),"done":x.get('done_reason')} for x in allres[-1]['full_speed']],
      "merge":allres[-1]['merge_precise']['grader']['pass'],
      "lru":allres[-1]['lru']['grader']['pass'],
      "calibration":allres[-1]['calibration']['pass']},indent=2),flush=True)
open('addendum.json','w').write(json.dumps(allres,indent=2))