"""Export selected local checkpoint into a native inference bundle."""
import hashlib
import json
import shutil
import sys
from pathlib import Path
sys.path.insert(0,str(Path(__file__).resolve().parents[1]))
import numpy as np
import onnxruntime as ort
import torch
from transformers import AutoModelForTokenClassification, AutoTokenizer
from pii.core import LABELS, MODEL_REVISION

class Export(torch.nn.Module):
    def __init__(self,model):
        super().__init__();self.model=model
    def forward(self,input_ids,attention_mask):
        return self.model(input_ids=input_ids,attention_mask=attention_mask).logits

def main():
    out=Path("artifacts/bundle");out.mkdir(parents=True,exist_ok=True)
    training=json.loads(Path("artifacts/training/run-config.json").read_text())
    source=Path("artifacts/training/best")
    tokenizer=AutoTokenizer.from_pretrained(source,local_files_only=True)
    model=AutoModelForTokenClassification.from_pretrained(source,local_files_only=True,attn_implementation="eager").eval()
    sample=tokenizer("🙂 José Smith: jose@example.com",return_tensors="pt")
    tokenizer.save_pretrained(out)
    config={"labels":LABELS,"max_length":training["max_length"],"stride":training["stride"],"model_revision":MODEL_REVISION,"dataset_revision":"b70ffaf5ff39e079776134c5bf4381f00a9fd1ed","offsets":"half-open UTF-8 bytes","raw_score":"minimum predicted-class token softmax within span","status":"experimental prototype, uncalibrated until calibration_file supplied"}
    (out/"config.json").write_text(json.dumps(config,indent=2)+"\n")
    wrapper=Export(model).eval()
    torch.onnx.export(wrapper, (sample["input_ids"],sample["attention_mask"]),str(out/"model.onnx"),input_names=["input_ids","attention_mask"],output_names=["logits"],dynamic_axes={"input_ids":{0:"batch",1:"sequence"},"attention_mask":{0:"batch",1:"sequence"},"logits":{0:"batch",1:"sequence"}},opset_version=17,dynamo=False)
    assert not model.training, "Export must preserve evaluation mode for parity checks"
    dylibs=list((Path(ort.__file__).parent/"capi").glob("libonnxruntime.*.dylib"))
    if len(dylibs)!=1:raise RuntimeError(f"Ambiguous native runtime paths: {dylibs}")
    shutil.copy2(dylibs[0],out/"libonnxruntime.dylib")
    shutil.copytree("rust/licenses",out/"licenses",dirs_exist_ok=True)
    shutil.copy2("artifacts/backbone/README.md",out/"licenses/backbone-model-card.md")
    (out/"licenses/dataset-attribution.txt").write_text("Nemotron-PII by NVIDIA, CC BY 4.0. Source: https://huggingface.co/datasets/nvidia/Nemotron-PII\nRevision b70ffaf5ff39e079776134c5bf4381f00a9fd1ed. Training filters and preparation modifications: see artifacts/data-manifest.json and docs/data.md.\n")
    options=ort.SessionOptions();options.intra_op_num_threads=1
    session=ort.InferenceSession(str(out/"model.onnx"),sess_options=options,providers=["CPUExecutionProvider"])
    checks=[]
    for texts in ([""],["Alice Smith", "🙂 José García: jose@example.com"],["Long ticket. "*170]):
        x=tokenizer(texts,return_tensors="pt",padding=True,truncation=True,max_length=training["max_length"])
        x={k:v for k,v in x.items() if k in ("input_ids","attention_mask")}
        with torch.no_grad():expected=model(**x).logits.numpy()
        actual=session.run(None,{k:v.numpy() for k,v in x.items()})[0]
        checks.append({"batch":len(texts),"sequence":x["input_ids"].shape[1],"max_abs_difference":float(np.abs(expected-actual).max()),"argmax_identical":bool(np.array_equal(expected.argmax(-1),actual.argmax(-1)))})
        if not checks[-1]["argmax_identical"]:raise RuntimeError("Export changed probe token predictions")
    (out/"export-evidence.json").write_text(json.dumps(checks,indent=2)+"\n")
    print(json.dumps(checks,indent=2),flush=True)

if __name__=="__main__":main()