"""Box-side deterministic item builder (no models). Implements METRIC_DEFINITIONS.json v1."""
import json, re, hashlib, sys
import pyarrow.parquet as pq
from tokenizers import Tokenizer
B='build/'; P='package/'
TOK={'STASIS_12269':Tokenizer.from_file(P+'stasis_tokenizer/tokenizer.json'),'GPT2_SMALL':Tokenizer.from_file(P+'gpt2/tokenizer.json')}
BOS={'STASIS_12269':2,'GPT2_SMALL':50256}
def byte_decoder():
    values=list(range(ord("!"),ord("~")+1))+list(range(161,173))+list(range(174,256)); chars=values[:]; e=0
    for i in range(256):
        if i not in values: values.append(i); chars.append(256+e); e+=1
    return {chr(c):v for v,c in zip(values,chars)}
BD=byte_decoder()
def enc(m,raw):
    if not raw: return []
    ids=TOK[m].encode(raw.decode('utf-8'),add_special_tokens=False).ids
    back=bytes(BD[c] for t in ids for c in TOK[m].id_to_token(t))
    assert back==raw,(m,raw[:50]); return ids
def utf8_prefix(raw,cap):
    c=raw[:cap]
    while c:
        try: c.decode('utf-8'); return c
        except UnicodeDecodeError as e: c=c[:e.start]
    return b''
def utf8_suffix(raw,cap):
    s=max(0,len(raw)-cap)
    while s<len(raw) and raw[s]&0xC0==0x80: s+=1
    return raw[s:]
def pair(ctx,tgt):
    out={}
    for m in TOK:
        c,t=enc(m,ctx),enc(m,tgt); out[m]={'ctx':c,'tgt':t}
    return out
def fits(p): return all(1+len(v['ctx'])+len(v['tgt'])<=1024 for v in p.values())
def wjson(name,rows):
    with open(P+'data/'+name,'w') as f:
        for r in rows: f.write(json.dumps(r,separators=(',',':'))+'\n')
    print(name,len(rows))
# --- WT2 full
rows=[r['text'] for r in pq.read_table(B+'data/wikitext2_raw_test.parquet').to_pylist()]
arts=[[]]
for t in rows:
    if re.match(r'^ = [^=].* = \n$',t): arts.append([])
    arts[-1].append(t)
items=[]
for ai,a in enumerate(arts):
    raw=''.join(a).encode()
    if not raw: continue
    off=0
    while off<len(raw):
        tgt=utf8_prefix(raw[off:],1024); assert tgt
        for cap in [2048,1536,1024,512,256,0]:
            ctx=utf8_suffix(raw[:off],cap); p=pair(ctx,tgt)
            if fits(p): break
        else: raise SystemExit('no fit')
        items.append({'unit':f'article/{ai}','start':off,'tgt_bytes':len(tgt),'ctx_bytes':len(ctx),'ctx_cap':cap,'ids':p}); off+=len(tgt)
    assert off==len(raw)
wjson('WT2_FULL_BPB.jsonl',items)
# --- panel continuity
panel=[json.loads(l) for l in open('/tmp/g4/runtime/packet/english.jsonl')]
items=[]
for r in panel:
    raw=utf8_prefix(r['text'].encode(),2048); off=0
    if not raw: items.append({'unit':r['id'],'start':0,'tgt_bytes':0,'ctx_bytes':0,'ids':None}); continue
    while off<len(raw):
        tgt=utf8_prefix(raw[off:],256); ctx=utf8_suffix(raw[:off],512); p=pair(ctx,tgt); assert fits(p)
        items.append({'unit':r['id'],'start':off,'tgt_bytes':len(tgt),'ctx_bytes':len(ctx),'ids':p}); off+=len(tgt)
wjson('WT2_PANEL256_CONTINUITY.jsonl',items)
# --- LAMBADA
items=[]
for i,r in enumerate(pq.read_table(B+'data/lambada_openai_test.parquet').to_pylist()):
    c,w=r['text'].rsplit(' ',1); tgt=(' '+w).encode(); p=pair(c.encode(),tgt); assert fits(p)
    items.append({'unit':f'lambada/{i}','tgt_bytes':len(tgt),'ids':p})
wjson('LAMBADA_OPENAI.jsonl',items)
# --- MC
trims=0
def mc(name,uid,ctx,choices,gold):
    global trims
    raw=ctx.encode(); chs=[]
    for ch in choices:
        t=(' '+ch).encode(); cap=len(raw)
        while True:
            p=pair(utf8_suffix(raw,cap),t)
            if fits(p): break
            trims+=1; cap=cap*3//4
        chs.append({'tgt_bytes':len(t),'ids':p})
    return {'unit':uid,'gold':gold,'choices':chs}
def hs_pre(t):
    t=t.strip().replace(' [title]','. '); t=re.sub('\\[.*?\\]','',t); return t.replace('  ',' ')
items=[mc('HS',f"hellaswag/{i}",hs_pre(r['activity_label']+': '+r['ctx_a']+' '+r['ctx_b'].capitalize()),[hs_pre(e) for e in r['endings']],int(r['label'])) for i,r in enumerate(pq.read_table(B+'data/hellaswag_validation.parquet').to_pylist())]
wjson('HELLASWAG.jsonl',items)
items=[mc('PIQA',f"piqa/{i}",'Question: '+r['goal']+'\nAnswer:',[r['sol1'],r['sol2']],int(r['label'])) for i,r in enumerate(pq.read_table(B+'data/piqa_validation.parquet').to_pylist())]
wjson('PIQA.jsonl',items)
items=[mc('ARC',r['id'],'Question: '+r['question']+'\nAnswer:',list(r['choices']['text']),list(r['choices']['label']).index(r['answerKey'])) for r in pq.read_table(B+'data/arc_easy_test.parquet').to_pylist()]
wjson('ARC_EASY.jsonl',items)
print('mc_context_trims',trims)
json.dump({'mc_context_trims':trims},open(P+'data/BUILD_STATS.json','w'))