"""Prefix dataset for early-kill: each run cut at each completed edit k, with the elapsed-time
fraction at that edit and the workspace reconstructed through edits[:k].
→ data/prefixes.jsonl  {id, task, split, y, k, n_edits, t_frac, remaining_s, total_s, files}"""
import json, os, glob, re
from datetime import datetime
ROOT=os.path.dirname(os.path.dirname(os.path.abspath(__file__))); SESS=os.path.expanduser('~/.pillbox/global/sessions'); PB=os.path.expanduser('~/code/pillbox/scripts/eval')
def ts(s):
    s=s.replace('Z','+00:00'); m=re.match(r'(.*\.)(\d+)(\+00:00)$',s)
    if m: s=m.group(1)+m.group(2)[:6].ljust(6,'0')+m.group(3)
    return datetime.fromisoformat(s)
tasks={}
for pat in ['tasks/*','sensitivity-tasks/*','toolz-tasks/*']:
    for d in glob.glob(os.path.join(PB,pat)):
        if not os.path.isdir(d): continue
        ws=os.path.join(d,'workspace'); files={}
        for root,_,fs in os.walk(ws):
            for f in fs:
                if not f.startswith('.'):
                    try: files[os.path.relpath(os.path.join(root,f),ws)]=open(os.path.join(root,f),errors='replace').read()
                    except: pass
        tasks[os.path.basename(d)]=files
EDIT={'edit','write','Edit','Write','MultiEdit','apply_patch','write_file','create_file'}
out=open(os.path.join(ROOT,'data/prefixes.jsonl'),'w'); n_rows=n_pre=0
for line in open(os.path.join(ROOT,'data/heldout.jsonl')):
    r=json.loads(line); p=os.path.join(SESS,r['id'],'log.jsonl')
    if not os.path.exists(p): continue
    evs=[json.loads(l) for l in open(p,errors='replace')]
    if len(evs)<2: continue
    t0=ts(evs[0]['at']); t_end=ts(evs[-1]['at']); total=(t_end-t0).total_seconds()
    if total<=0: continue
    edits=[e for e in evs if e['payload'].get('type')=='tool_call' and e['payload'].get('status')=='completed' and str(e['payload'].get('name')) in EDIT and e['payload'].get('input')]
    if not edits: continue
    ws=dict(tasks.get(r['task'],{})); touched=set(); n_rows+=1
    for k,e in enumerate(edits,1):
        i=e['payload']['input']; name=(i.get('filePath') or i.get('path') or i.get('file_path') or '').replace('/workspace/','').split('/',1)[-1] if (i.get('filePath') or i.get('path') or i.get('file_path')) else ''
        base=os.path.basename(name); key=name if name in ws else next((kk for kk in ws if os.path.basename(kk)==base),None) or name
        old=i.get('oldString') or i.get('old_string') or ''; new=i.get('newString') or i.get('new_string') or i.get('content') or ''
        if not old or key not in ws or old not in ws.get(key,''): ws[key]=(ws.get(key,'')+'\n'+new) if old else new
        else: ws[key]=ws[key].replace(old,new,1)
        touched.add(base)
        t=(ts(e['at'])-t0).total_seconds()
        files={kk:v[:12000] for kk,v in ws.items() if os.path.basename(kk) in touched and 'test' not in os.path.basename(kk).lower()}
        out.write(json.dumps({'id':r['id'],'task':r['task'],'split':r['split'],'y':1 if r['label']['passed'] else 0,'k':k,'n_edits':len(edits),'t_frac':round(t/total,3),'remaining_s':round(total-t,1),'total_s':round(total,1),'prompt':r['prompt'][:5000],'files':files})+'\n'); n_pre+=1
print(f'rows {n_rows} prefixes {n_pre}')