"""Oracle BIO logits prove window mechanics independently of trained predictions."""
import json
import sys
from pathlib import Path
import numpy as np
from transformers import AutoTokenizer
sys.path.insert(0,str(Path(__file__).resolve().parents[1]))
from pii.core import encode_windows, LABELS, byte_positions, decode_window, result

root=Path(__file__).resolve().parents[1]
tokenizer=AutoTokenizer.from_pretrained(root/'artifacts/backbone',local_files_only=True)
max_length,stride=512,42
cases=[]

def make_case(name,prefix,entities,suffix):
    text=prefix; spans=[]
    for label,value in entities:
        start=len(text);text+=value;spans.append({'label':label,'start':start,'end':len(text)});text+=' '
    text+=suffix
    bp=byte_positions(text)
    expected=[{'label':s['label'],'start':bp[s['start']],'end':bp[s['end']]} for s in spans]
    encoded=encode_windows(tokenizer,text,max_length,stride)
    raw=tokenizer(text,add_special_tokens=False,truncation=False,return_offsets_mapping=True,verbose=False)
    covered=[]
    for wi,ids in enumerate(encoded['input_ids']):covered.extend(ids[1 if wi==0 else stride+1:-1])
    assert covered==raw['input_ids'], 'manual windows lost or duplicated tokens outside overlap'
    windows=[]; findings=[]; covers=[0]*len(spans); clips=0
    for wi,offsets in enumerate(encoded['offset_mapping']):
        content=[(a,b) for a,b in offsets if a!=b]
        complete=[content[0][0]<=s['start'] and content[-1][1]>=s['end'] for s in spans]
        covers[:]=[v+int(c) for v,c in zip(covers,complete)]
        logits=np.full((len(offsets),15),-8,dtype=np.float32)
        for i,(a,b) in enumerate(offsets):
            label_id=0; peak=12.
            if a!=b:
                for si,s in enumerate(spans):
                    if a<s['end'] and s['start']<b:
                        assert s['start']<=a and b<=s['end'], 'fixture boundary unrepresentable'
                        label_id=1+2*LABELS.index(s['label'])+int(a!=s['start'])
                        peak=5. if complete[si] else 12.
                        break
            logits[i,label_id]=peak
        decoded=decode_window(logits,offsets,text,max_length,wi>0,wi+1<len(encoded['input_ids']))
        clips+=sum(f.get('_boundary_fragment',False) for f in decoded)
        findings.extend(decoded)
        windows.append({'input_ids':encoded['input_ids'][wi],'offsets':[[bp[a],bp[b]] for a,b in offsets], 'logits':logits.tolist()})
    actual=result(text,findings)
    assert [{k:f[k] for k in ('label','start','end')} for f in actual['findings']]==expected,(name,actual,expected)
    assert all(covers), 'no complete window for entity'
    # Demonstrate the old greedy rule really loses exact spans on boundary cases.
    old=result(text,[{k:v for k,v in f.items() if k!='_boundary_fragment'} for f in findings])
    old_wrong=[{k:f[k] for k in ('label','start','end')} for f in old['findings']]!=expected
    if name!='adjacent-repeated-unicode':assert clips and old_wrong
    cases.append({'name':name,'text':text,'expected':expected,'redacted':actual['redacted'],'windows':windows,'complete_window_counts':covers,'edge_predictions':clips,'old_greedy_would_fail':old_wrong,'full_token_coverage':True,'raw_tokens':len(raw['input_ids'])})

make_case('boundary-email','🙂 '+'hello '*505,[('email','jose.garcia@example.com')],'end '*30)
make_case('max43-token-boundary-entity','🙂 '+'hello '*490,[('street_address',' '.join(['alpha']*43))],'end '*30)
make_case('adjacent-repeated-unicode','🙂 ',[('first_name','José'),('last_name','García'),('first_name','José'),('last_name','García')],'done')
fixture={'max_length':max_length,'stride':stride,'labels':LABELS,'cases':cases}
(root/'rust/tests/boundary-oracle.json').write_text(json.dumps(fixture,separators=(',',':'))+'\n')
report={'max_length':max_length,'stride':stride,'policy':'Artificial-edge fragments suppressed only when a strictly containing same-label prediction exists. Remaining candidates use original greedy score ordering.', 'proof':'Oracle true BIO labels; partial-window probabilities intentionally exceed complete-window probabilities.', 'cases':[{k:v for k,v in c.items() if k not in ('text','windows','redacted')} for c in cases]}
(root/'artifacts/boundary-proof.json').write_text(json.dumps(report,indent=2)+'\n')
print(json.dumps(report,indent=2))