import json, os, sys, time, urllib.request
KEY=open('/home/box/.secrets/openrouter_api_key').read().strip()
SYSTEM=open('repo/skills/ponytail/SKILL.md').read()
TASKS={
 'T1':"Write me a Python function that validates email addresses.",
 'T2':"Write a reusable debounce function in vanilla JavaScript: debounce(fn, delay) returns a debounced version of fn that delays calling it until delay ms after the last call.",
 'T3':"Write Python code that reads sales.csv and sums the 'amount' column.",
 'T4':"Write a Python function save_upload(filename, data) that saves user-uploaded bytes into the ./uploads directory under the given filename.",
 'T5':"Write a Python function parse_price(s) that converts a price string like '$1,234.50' into integer cents.",
}
model, runs = sys.argv[1], int(sys.argv[2])
tag = model.split('/')[1].split(':')[0]
os.makedirs('out', exist_ok=True)
for r in range(1, runs+1):
  for t, task in TASKS.items():
    for arm in ('baseline','ponytail'):
      fn=f'out/{tag}_{t}_{arm}_r{r}.json'
      if os.path.exists(fn): continue
      msgs=([{'role':'system','content':SYSTEM}] if arm=='ponytail' else [])+[{'role':'user','content':task}]
      body=json.dumps({'model':model,'messages':msgs,'max_tokens':8192}).encode()
      req=urllib.request.Request('https://openrouter.ai/api/v1/chat/completions',body,{'Authorization':'Bearer '+KEY,'Content-Type':'application/json'})
      t0=time.time()
      try:
        d=json.load(urllib.request.urlopen(req,timeout=300)); err=None
      except urllib.error.HTTPError as e:
        d=None; err=f'HTTP {e.code}: {e.read()[:300]!r}'
      except Exception as e:
        d=None; err=repr(e)[:300]
      dt=round(time.time()-t0,1)
      if d and 'choices' in d:
        rec={'model':model,'task':t,'arm':arm,'run':r,'secs':dt,'usage':d.get('usage'),'content':d['choices'][0]['message'].get('content'),'finish':d['choices'][0].get('finish_reason'),'provider':d.get('provider')}
        json.dump(rec,open(fn,'w'),indent=1)
        print(time.strftime('%H:%M:%S'),fn,dt,'s',(d.get('usage') or {}).get('completion_tokens'),flush=True)
      else:
        print(time.strftime('%H:%M:%S'),fn,'ERROR',err or str(d)[:300],flush=True)
        if err and '429' in err: time.sleep(20)