"""Rebuild each row's final workspace files (task workspace + agent edits applied in order)
and collect the visible test files. Output: data/files/<id>.json  {files:{path:content}, tests:{path:content}, applied, unapplied}"""
import json, os, glob, sys
PB=os.path.expanduser('~/code/pillbox/scripts/eval')
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 f.startswith('.'): continue
                p=os.path.join(root,f); rel=os.path.relpath(p,ws)
                try: files[rel]=open(p,errors='replace').read()
                except: pass
        tasks[os.path.basename(d)]=files
os.makedirs('data/files',exist_ok=True)
stats={'rows':0,'clean':0,'partial':0,'no_ws_match':0}
for line in open('data/heldout.jsonl'):
    r=json.loads(line); ws=dict(tasks.get(r['task'],{}))
    applied=unapplied=0
    for e in r['edits']:
        name=e['file'].replace('/ws/','').lstrip('/'); base=os.path.basename(name)
        key=name if name in ws else next((k for k in ws if os.path.basename(k)==base),None) or name
        if e['tool'] in ('write','Write','write_file','create_file') or (not e['old'] and key not in ws):
            ws[key]=e['new']; applied+=1
        elif e['old'] and key in ws and e['old'] in ws[key]:
            ws[key]=ws[key].replace(e['old'],e['new'],1); applied+=1
        else:
            unapplied+=1
            ws.setdefault(key,''); ws[key]=ws[key]+'\n'+e['new']  # best effort: append the new text
    touched={os.path.basename(x['file']) for x in r['edits']}
    files={k:v[:20000] for k,v in ws.items() if os.path.basename(k) in touched and not k.endswith('_test.py') and 'test' not in os.path.basename(k).lower()}
    tests={k:v[:12000] for k,v in ws.items() if k.endswith('_test.py') or os.path.basename(k).startswith('test_') or '/test' in k}
    json.dump({'id':r['id'],'task':r['task'],'files':files,'tests':tests,'applied':applied,'unapplied':unapplied},open(f"data/files/{r['id']}.json",'w'))
    stats['rows']+=1; stats['clean' if unapplied==0 else 'partial']+=1
    if not files: stats['no_ws_match']+=1
print(stats)