"""Add CPU-oriented INT8 and INT4 variants to an exported artifact bundle."""

from __future__ import annotations

import argparse
import json
import logging
import os
import shutil
import tempfile
from pathlib import Path

import onnxruntime as ort
from onnxruntime.quantization import QuantType, quantize_dynamic
from onnxruntime.quantization.matmul_nbits_quantizer import (
    DefaultWeightOnlyQuantConfig,
    MatMulNBitsQuantizer,
)

from .artifact import variant_metadata, write_manifest


DECISION_HEAD = "/openthai/head/SlotMatMul"
EXPECTED_BACKBONE_MATMULS = 186


def _assert_runtime_contract(path: Path) -> None:
    session = ort.InferenceSession(str(path), providers=["CPUExecutionProvider"])
    expected_inputs = {
        "input_ids",
        "attention_mask",
        "answer_positions",
        "option_counts",
        "qtypes",
    }
    if {value.name for value in session.get_inputs()} != expected_inputs:
        raise RuntimeError(f"quantized graph has an unexpected input contract: {path}")
    if {value.name for value in session.get_outputs()} != {"logits", "probabilities"}:
        raise RuntimeError(f"quantized graph has an unexpected output contract: {path}")


def _assert_operator_contract(metadata: dict, precision: str) -> None:
    counts = metadata["operator_counts"]
    expected = "ai.onnx::MatMulInteger" if precision == "int8" else "com.microsoft::MatMulNBits"
    if counts.get(expected) != EXPECTED_BACKBONE_MATMULS:
        raise RuntimeError(
            f"{precision} graph must contain {EXPECTED_BACKBONE_MATMULS} {expected} nodes; "
            f"found {counts.get(expected, 0)}"
        )
    if counts.get("ai.onnx::MatMul") != 1:
        raise RuntimeError(f"{precision} graph must retain exactly one FP32 decision-head MatMul")


def _build_int8(source: Path, destination: Path) -> None:
    quantize_dynamic(
        source,
        destination,
        op_types_to_quantize=["MatMul"],
        per_channel=True,
        weight_type=QuantType.QInt8,
        nodes_to_exclude=[DECISION_HEAD],
        use_external_data_format=True,
    )


def _build_int4(source: Path, destination: Path) -> None:
    config = DefaultWeightOnlyQuantConfig(
        bits=4,
        block_size=32,
        is_symmetric=False,
        op_types_to_quantize=("MatMul",),
    )
    quantizer = MatMulNBitsQuantizer(
        str(source),
        nodes_to_exclude=[DECISION_HEAD],
        algo_config=config,
    )
    quantizer.process()
    quantizer.model.save_model_to_file(str(destination), True)


def _data_path(graph: Path) -> Path:
    return Path(f"{graph}.data")


def _refresh_release_files(artifact: Path) -> None:
    """Refresh packaged client/tooling and documentation from this checkout."""
    project_root = Path(__file__).resolve().parents[2]
    for source_name, destination_name in (
        ("LICENSE", "LICENSE"),
        ("MODEL_CARD.md", "README.md"),
        ("MODEL_CARD.md", "MODEL_CARD.md"),
        ("ATTRIBUTION.md", "ATTRIBUTION.md"),
        ("INSTRUCTIONS.md", "INSTRUCTIONS.md"),
        ("pyproject.toml", "pyproject.toml"),
        ("uv.lock", "uv.lock"),
    ):
        shutil.copy2(project_root / source_name, artifact / destination_name)
    for package in ("openthai_systemone_onnx", "openthai_systemone_onnx_tools"):
        shutil.copytree(
            project_root / "src" / package,
            artifact / "src" / package,
            dirs_exist_ok=True,
            ignore=shutil.ignore_patterns("__pycache__", "*.pyc"),
        )
    shutil.copytree(
        project_root / "examples",
        artifact / "examples",
        dirs_exist_ok=True,
        ignore=shutil.ignore_patterns("__pycache__", "*.pyc"),
    )


def quantize_artifact(artifact: Path) -> Path:
    artifact = artifact.resolve()
    source = artifact / "onnx" / "model.onnx"
    manifest_path = artifact / "artifact-manifest.json"
    if not source.is_file() or not manifest_path.is_file():
        raise FileNotFoundError("artifact must contain onnx/model.onnx and artifact-manifest.json")

    onnx_dir = source.parent
    with tempfile.TemporaryDirectory(prefix=".quantize-", dir=onnx_dir) as temporary:
        temporary_dir = Path(temporary)
        temporary_graphs = {
            "int8": temporary_dir / "model.int8.onnx",
            "int4": temporary_dir / "model.int4.onnx",
        }
        previous_logging_disable = logging.root.manager.disable
        logging.disable(logging.INFO)
        try:
            _build_int8(source, temporary_graphs["int8"])
            _build_int4(source, temporary_graphs["int4"])
        finally:
            logging.disable(previous_logging_disable)

        temporary_variants = {}
        quantization = {
            "int8": {
                "mode": "dynamic",
                "weight_type": "qint8",
                "activation_quantization": "dynamic",
                "per_channel": True,
                "operator": "MatMulInteger",
                "excluded_nodes": [DECISION_HEAD],
            },
            "int4": {
                "mode": "weight_only",
                "weight_type": "uint4",
                "activation_type": "float32",
                "block_size": 32,
                "symmetric": False,
                "operator": "MatMulNBits",
                "excluded_nodes": [DECISION_HEAD],
            },
        }
        for precision, graph in temporary_graphs.items():
            if not _data_path(graph).is_file():
                raise RuntimeError(f"quantizer did not create external data for {precision}")
            _assert_runtime_contract(graph)
            metadata = variant_metadata(
                temporary_dir,
                precision=precision,
                graph=graph.name,
                external_data=_data_path(graph).name,
                quantization=quantization[precision],
            )
            _assert_operator_contract(metadata, precision)
            temporary_variants[precision] = metadata

        for precision, graph in temporary_graphs.items():
            final_graph = onnx_dir / graph.name
            final_data = _data_path(final_graph)
            os.replace(_data_path(graph), final_data)
            os.replace(graph, final_graph)

    old_manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
    variants = dict(old_manifest.get("variants", {}))
    if "fp32" not in variants:
        variants["fp32"] = variant_metadata(
            artifact,
            precision="fp32",
            graph="onnx/model.onnx",
            external_data="onnx/model.onnx.data",
            quantization=None,
        )
    for precision, metadata in temporary_variants.items():
        metadata["graph"] = f"onnx/{metadata['graph']}"
        metadata["external_data"] = f"onnx/{metadata['external_data']}"
        variants[precision] = metadata

    _refresh_release_files(artifact)
    report = artifact / "validation-report.json"
    report.unlink(missing_ok=True)
    return write_manifest(artifact, extra={"variants": variants})


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--artifact", type=Path, required=True)
    return parser


def main(argv: list[str] | None = None) -> None:
    args = build_parser().parse_args(argv)
    manifest = quantize_artifact(args.artifact)
    print(f"Quantized bundle complete: {manifest}")


if __name__ == "__main__":
    main()