"""Export the desklib DeBERTa-v3 AI-detector to a small int8 ONNX so it runs on
onnxruntime (no torch at inference) and frees ~1.3GB of disk.

Disk-safe ordering on a near-full disk: copy the tokenizer out, load the torch
model into RAM, DELETE the 1.7GB HF cache (model is already in RAM), THEN write
the ONNX — so the disk never holds the torch cache and the fp32 ONNX at once.
Finally verify the int8 ONNX matches torch on sample texts before we rely on it.

Output: models/desklib_onnx/{model.onnx, tokenizer files, config.json, meta.json}
Run: python scripts/export_desklib_onnx.py
"""
import glob, json, os, shutil, sys
import numpy as np

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
DST = os.path.join(ROOT, "models", "desklib_onnx")
NAME = "desklib/ai-text-detector-v1.01"
TOK_FILES = ["spm.model", "tokenizer.json", "tokenizer_config.json",
             "special_tokens_map.json", "added_tokens.json", "config.json"]


def main():
    import torch, torch.nn as nn
    from transformers import AutoConfig, AutoModel, AutoTokenizer, PreTrainedModel

    os.makedirs(DST, exist_ok=True)

    class Desklib(PreTrainedModel):
        config_class = AutoConfig

        def __init__(self, config):
            super().__init__(config)
            self.model = AutoModel.from_config(config)
            self.classifier = nn.Linear(config.hidden_size, 1)
            self.init_weights()

        def forward(self, input_ids, attention_mask=None, **kw):
            out = self.model(input_ids, attention_mask=attention_mask)[0]
            mask = attention_mask.unsqueeze(-1).expand(out.size()).float()
            pooled = (out * mask).sum(1) / mask.sum(1).clamp(min=1e-9)
            return self.classifier(pooled)

    tok = AutoTokenizer.from_pretrained(NAME)
    model = Desklib.from_pretrained(NAME).eval()
    print("torch model loaded into RAM")

    # copy tokenizer/config out of the (now-populated) cache before we delete it
    snap = glob.glob(os.path.join(
        os.path.expanduser("~"), ".cache", "huggingface", "hub",
        "models--desklib--ai-text-detector-v1.01", "snapshots", "*"))
    if snap:
        for f in TOK_FILES:
            src = os.path.join(snap[0], f)
            if os.path.exists(src):
                shutil.copy(src, os.path.join(DST, f))
        print("copied tokenizer/config to", DST)

    # sample texts for verification (computed BEFORE we delete anything)
    samples = ["The proposed framework leverages a comprehensive and multifaceted "
               "approach to optimize performance across diverse benchmarks, underscoring "
               "its pivotal role in advancing the field.",
               "we ran the test three times and honestly the numbers were all over the "
               "place, not sure why, maybe the sensor was loose or we messed up the wiring"]

    @torch.no_grad()
    def torch_p(t):
        e = tok(t, truncation=True, max_length=512, return_tensors="pt")
        return float(torch.sigmoid(model(**{k: e[k] for k in
                     ("input_ids", "attention_mask")})[0]))
    torch_ref = [torch_p(t) for t in samples]

    # free the 1.7GB disk cache now that the model is in RAM
    cache_dir = os.path.join(os.path.expanduser("~"), ".cache", "huggingface",
                             "hub", "models--desklib--ai-text-detector-v1.01")
    shutil.rmtree(cache_dir, ignore_errors=True)
    print("freed HF cache (model stays in RAM)")

    # export fp32 ONNX
    fp32 = os.path.join(DST, "model_fp32.onnx")
    dummy = tok(samples[0], return_tensors="pt", truncation=True, max_length=64)
    torch.onnx.export(
        model, (dummy["input_ids"], dummy["attention_mask"]), fp32,
        input_names=["input_ids", "attention_mask"], output_names=["logits"],
        dynamic_axes={"input_ids": {0: "b", 1: "s"},
                      "attention_mask": {0: "b", 1: "s"}, "logits": {0: "b"}},
        opset_version=14, do_constant_folding=True, dynamo=False)
    print(f"exported fp32 ONNX ({os.path.getsize(fp32)//1024//1024} MB)")

    # dynamic int8 quantization
    from onnxruntime.quantization import quantize_dynamic, QuantType
    out = os.path.join(DST, "model.onnx")
    quantize_dynamic(fp32, out, weight_type=QuantType.QInt8)
    os.remove(fp32)
    print(f"quantized int8 ONNX ({os.path.getsize(out)//1024//1024} MB); removed fp32")

    # verify ONNX matches torch
    import onnxruntime as ort
    sess = ort.InferenceSession(out, providers=["CPUExecutionProvider"])
    inn = [i.name for i in sess.get_inputs()]

    def onnx_p(t):
        e = tok(t, truncation=True, max_length=512)
        feed = {"input_ids": np.array([e["input_ids"]], np.int64),
                "attention_mask": np.array([e["attention_mask"]], np.int64)}
        logit = sess.run(None, {k: v for k, v in feed.items() if k in inn})[0]
        return float(1 / (1 + np.exp(-logit.reshape(-1)[0])))

    print("\nVERIFY (torch vs int8 ONNX):")
    ok = True
    for t, ref in zip(samples, torch_ref):
        o = onnx_p(t)
        d = abs(o - ref)
        ok &= d < 0.06
        print(f"  torch={ref:.3f}  onnx={o:.3f}  |Δ|={d:.3f}")
    json.dump({"model": NAME, "format": "onnx-int8-dynamic", "max_len": 512,
               "verified": ok}, open(os.path.join(DST, "meta.json"), "w"), indent=1)
    print("\nVERIFIED OK" if ok else "\nWARNING: ONNX deviates from torch — check")


if __name__ == "__main__":
    main()