#!/usr/bin/env python3
"""Export feyninc/pulpie-orange-small to optimized ONNX."""

from __future__ import annotations

import argparse
import json
import shutil
from collections import Counter
from pathlib import Path
from typing import Any

import onnx
import onnxruntime as ort
import torch
from onnxruntime.quantization import QuantType, quantize_dynamic
from transformers import AutoModelForTokenClassification, AutoTokenizer


TOKENIZER_FILES = (
    "config.json",
    "tokenizer.json",
    "tokenizer_config.json",
    "special_tokens_map.json",
    "configuration_eurobert.py",
    "modeling_eurobert.py",
)


class LogitsOnly(torch.nn.Module):
    def __init__(self, model: torch.nn.Module):
        super().__init__()
        self.model = model

    def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
        return self.model(input_ids=input_ids, attention_mask=attention_mask).logits


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--source-model", default="artifacts/source_model")
    parser.add_argument("--output-dir", default="artifacts/onnx_model")
    parser.add_argument("--report", default="artifacts/reports/export_report.json")
    parser.add_argument("--opset", type=int, default=17)
    parser.add_argument("--sample-seq-len", type=int, default=96)
    parser.add_argument("--no-quantize", action="store_true")
    return parser.parse_args()


def file_size(path: Path) -> int:
    return path.stat().st_size if path.exists() else 0


def graph_summary(path: Path) -> dict[str, Any]:
    model = onnx.load(str(path), load_external_data=False)
    domains = sorted({node.domain for node in model.graph.node if node.domain not in ("", "ai.onnx")})
    ops = Counter(node.op_type if not node.domain else f"{node.domain}::{node.op_type}" for node in model.graph.node)
    return {
        "ir_version": model.ir_version,
        "opsets": {op.domain or "ai.onnx": op.version for op in model.opset_import},
        "node_count": len(model.graph.node),
        "nonstandard_domains": domains,
        "top_ops": dict(ops.most_common(30)),
    }


def check_standard_model(path: Path) -> tuple[bool, str | None, dict[str, Any]]:
    try:
        onnx.checker.check_model(str(path))
        summary = graph_summary(path)
        if summary["nonstandard_domains"]:
            return False, f"nonstandard domains: {summary['nonstandard_domains']}", summary
        return True, None, summary
    except Exception as exc:  # noqa: BLE001
        return False, repr(exc), {}


def optimize_with_ort(raw_path: Path, out_path: Path, level: ort.GraphOptimizationLevel) -> None:
    session_options = ort.SessionOptions()
    session_options.graph_optimization_level = level
    session_options.optimized_model_filepath = str(out_path)
    ort.InferenceSession(str(raw_path), sess_options=session_options, providers=["CPUExecutionProvider"])


def build_sample_inputs(tokenizer: Any, seq_len: int) -> tuple[torch.Tensor, torch.Tensor]:
    sep_id = tokenizer.convert_tokens_to_ids("<|sep|>")
    text = '<article _item_id="0"><h1>Pulpie ONNX export</h1><p>Main content block.</p></article>'
    token_ids = tokenizer.encode(text, add_special_tokens=False)
    ids = [tokenizer.bos_token_id] + token_ids + [sep_id] + [tokenizer.eos_token_id]
    if len(ids) < seq_len:
        filler = tokenizer.encode(" More HTML text.", add_special_tokens=False)
        while len(ids) + len(filler) + 1 < seq_len:
            ids[-1:-1] = filler
    ids = ids[:seq_len]
    ids[-1] = tokenizer.eos_token_id
    input_ids = torch.tensor([ids], dtype=torch.long)
    attention_mask = torch.ones_like(input_ids)
    return input_ids, attention_mask


def copy_runtime_files(source_dir: Path, output_dir: Path) -> None:
    for name in TOKENIZER_FILES:
        src = source_dir / name
        if src.exists():
            shutil.copy2(src, output_dir / name)


def main() -> None:
    args = parse_args()
    source_dir = Path(args.source_model).resolve()
    output_dir = Path(args.output_dir).resolve()
    report_path = Path(args.report).resolve()
    output_dir.mkdir(parents=True, exist_ok=True)
    report_path.parent.mkdir(parents=True, exist_ok=True)

    tokenizer = AutoTokenizer.from_pretrained(str(source_dir), trust_remote_code=True)
    model = AutoModelForTokenClassification.from_pretrained(
        str(source_dir),
        trust_remote_code=True,
        num_labels=2,
        dtype=torch.float32,
        attn_implementation="eager",
    ).eval()
    model.config.return_dict = True
    model.config.use_cache = False
    wrapped = LogitsOnly(model).eval()

    input_ids, attention_mask = build_sample_inputs(tokenizer, args.sample_seq_len)
    with torch.no_grad():
        logits = wrapped(input_ids, attention_mask)

    raw_path = output_dir / "model_raw.onnx"
    torch.onnx.export(
        wrapped,
        (input_ids, attention_mask),
        str(raw_path),
        input_names=["input_ids", "attention_mask"],
        output_names=["logits"],
        dynamic_axes={
            "input_ids": {0: "batch_size", 1: "sequence_length"},
            "attention_mask": {0: "batch_size", 1: "sequence_length"},
            "logits": {0: "batch_size", 1: "sequence_length"},
        },
        opset_version=args.opset,
        do_constant_folding=True,
        export_params=True,
    )

    optimizer_attempts: list[dict[str, Any]] = []
    selected_path: Path | None = None
    for level_name, level in (
        ("ORT_ENABLE_EXTENDED", ort.GraphOptimizationLevel.ORT_ENABLE_EXTENDED),
        ("ORT_ENABLE_BASIC", ort.GraphOptimizationLevel.ORT_ENABLE_BASIC),
    ):
        candidate = output_dir / f"model_{level_name.lower()}.onnx"
        try:
            optimize_with_ort(raw_path, candidate, level)
            ok, error, summary = check_standard_model(candidate)
            optimizer_attempts.append(
                {"level": level_name, "path": str(candidate), "standard_onnx": ok, "error": error, "summary": summary}
            )
            if ok:
                selected_path = candidate
                break
        except Exception as exc:  # noqa: BLE001
            optimizer_attempts.append(
                {"level": level_name, "path": str(candidate), "standard_onnx": False, "error": repr(exc)}
            )

    raw_ok, raw_error, raw_summary = check_standard_model(raw_path)
    if selected_path is None:
        if not raw_ok:
            raise RuntimeError(f"Raw ONNX failed checker: {raw_error}")
        selected_path = raw_path

    model_path = output_dir / "model.onnx"
    embedded_model = onnx.load(str(selected_path), load_external_data=True)
    onnx.save_model(embedded_model, str(model_path), save_as_external_data=False)

    final_ok, final_error, final_summary = check_standard_model(model_path)
    if not final_ok:
        raise RuntimeError(f"Final ONNX failed checker or standard-domain check: {final_error}")

    quantization: dict[str, Any] = {"enabled": not args.no_quantize}
    quantized_path = output_dir / "model_quantized.onnx"
    if not args.no_quantize:
        try:
            quantize_dynamic(
                str(model_path),
                str(quantized_path),
                weight_type=QuantType.QInt8,
                per_channel=True,
            )
            q_ok, q_error, q_summary = check_standard_model(quantized_path)
            quantization.update(
                {
                    "path": str(quantized_path),
                    "checker_passed": q_ok,
                    "error": q_error,
                    "summary": q_summary,
                    "size_bytes": file_size(quantized_path),
                }
            )
        except Exception as exc:  # noqa: BLE001
            quantization.update({"checker_passed": False, "error": repr(exc)})

    copy_runtime_files(source_dir, output_dir)

    report = {
        "source_model": str(source_dir),
        "output_dir": str(output_dir),
        "opset": args.opset,
        "sample_input_shape": list(input_ids.shape),
        "sample_logits_shape": list(logits.shape),
        "sample_logits_dtype": str(logits.dtype),
        "raw_model": {
            "path": str(raw_path),
            "checker_passed": raw_ok,
            "error": raw_error,
            "summary": raw_summary,
            "size_bytes": file_size(raw_path),
        },
        "optimizer_attempts": optimizer_attempts,
        "model": {
            "path": str(model_path),
            "checker_passed": final_ok,
            "summary": final_summary,
            "size_bytes": file_size(model_path),
        },
        "quantization": quantization,
    }
    report_path.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()