"""Check native Rust IDs, masks, offsets, logits, decoder/redaction and explicit errors."""
import argparse
import json
import subprocess
import sys
import time
from pathlib import Path
import numpy as np
import onnxruntime as ort
from transformers import AutoTokenizer
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from pii.core import encode_windows, byte_positions, decode_window, result, Predictor

def main():
    p=argparse.ArgumentParser()
    p.add_argument('bundle',type=Path)
    p.add_argument('--binary',type=Path,default=Path('rust/target/release/datafog-pii'))
    p.add_argument('--output',type=Path)
    args=p.parse_args()
    config=json.loads((args.bundle/'config.json').read_text())
    tokenizer=AutoTokenizer.from_pretrained(args.bundle,local_files_only=True)
    opts=ort.SessionOptions(); opts.intra_op_num_threads=1
    session=ort.InferenceSession(str(args.bundle/'model.onnx'),sess_options=opts,providers=['CPUExecutionProvider'])
    predictor=Predictor(args.bundle)
    calibration=json.loads((args.bundle/config['calibration_file']).read_text()) if config.get('calibration_file') else None
    texts=['','Alice Smith','🙂 José García: jose@example.com\nAccount: ABC123','Email marketing is useful.','Ann Ann, phone 415-555-0123.','Customer message. '*550+'\nJosé García <jose@example.com>','Order delay. '*170+'é Smith\n'+'More detail. '*180]
    payload=''.join(json.dumps({'text':text})+'\n' for text in texts)+'{"wrong":1}\n'+json.dumps({'text':'Recovery after invalid input'})+'\n'
    start=time.perf_counter()
    proc=subprocess.run([str(args.binary.resolve()),str(args.bundle.resolve()),'--debug'],input=payload,text=True,capture_output=True,check=True)
    elapsed=time.perf_counter()-start
    outputs=[json.loads(line) for line in proc.stdout.splitlines()]
    assert len(outputs)==len(texts)+2
    assert 'error' in outputs[-2] and 'contains_pii' not in outputs[-2],outputs[-2]
    assert 'error' not in outputs[-1],outputs[-1]
    max_diff=0.; checks=[]
    for text,out in zip(texts,outputs):
        assert 'error' not in out,out
        encoded=encode_windows(tokenizer,text,config['max_length'],config['stride'])
        assert len(encoded['input_ids'])==len(out['windows'])
        findings=[]
        for i,window in enumerate(out.pop('windows')):
            assert window['input_ids']==encoded['input_ids'][i]
            assert window['attention_mask']==encoded['attention_mask'][i]
            bp=byte_positions(text)
            assert window['offsets']==[[bp[a],bp[b]] for a,b in encoded['offset_mapping'][i]]
            logits=session.run(None,{k:np.array([encoded[k][i]],dtype=np.int64) for k in ('input_ids','attention_mask')})[0][0]
            rust_logits=np.array(window['logits'],dtype=np.float32)
            diff=float(np.abs(logits-rust_logits).max()); max_diff=max(max_diff,diff)
            # Same CPU runtime and one-thread settings: require exact f32 values.
            np.testing.assert_array_equal(logits,rust_logits)
            findings.extend(decode_window(logits,encoded['offset_mapping'][i],text,config['max_length'], i > 0, i + 1 < len(encoded['input_ids'])))
        expected=result(text,findings,calibration,calibration.get('threshold') if calibration else None)
        public=predictor.predict(text)
        assert public==expected, 'actual Predictor differs from shared reference'
        assert out['contains_pii']==expected['contains_pii']
        assert out['redacted']==expected['redacted']
        assert len(out['findings'])==len(expected['findings'])
        for actual,want in zip(out['findings'],expected['findings']):
            for key in ('start','end','label'):assert actual[key]==want[key]
            # f64 scalar summation vs numpy reduction differ at rounding scale.
            assert abs(actual['raw_score']-want['raw_score'])<=64*np.finfo(np.float64).eps
            if want['confidence'] is None: assert actual['confidence'] is None
            else: assert abs(actual['confidence']-want['confidence'])<=1e-10
        checks.append({'bytes':len(text.encode()),'windows':len(encoded['input_ids']),'findings':len(out['findings'])})
    report={'bundle':str(args.bundle),'cases':checks,'max_abs_logit_difference':max_diff,'logits_bit_identical':True,'token_ids_masks_global_byte_offsets_parity':True,'findings_redaction_parity':True,'actual_python_predictor_parity':True,'malformed_request_explicit_error_and_recovery':True,'cold_process_debug_batch_seconds':elapsed,'note':'Parity fixtures prove implementation consistency, not detection quality.'}
    serialized=json.dumps(report,indent=2)+'\n'
    if args.output:args.output.write_text(serialized)
    print(serialized)
if __name__=='__main__':main()