#!/usr/bin/env python3
"""Inspect an ONNX model and emit a machine-readable characterization.

Reports the facts a production-readiness review needs about a checkpoint before and
after optimization: opset, node/initializer counts, the weight-dtype histogram, whether
external data is used, the exact I/O tensor contract (names, dtypes, shapes), and the
operator histogram. Pure `onnx` (no onnxruntime) so it also runs on graphs a given ORT
build cannot execute.

Usage:
    python inspect_onnx.py MODEL.onnx [--json OUT.json]
"""

from __future__ import annotations

import argparse
import collections
import hashlib
import json
import os
import sys

import onnx
from onnx import TensorProto

# TensorProto enum value -> human dtype name, for the weight/IO histograms.
DTYPE_NAME = {v: k for k, v in TensorProto.DataType.items()}


def _elem_type_name(elem_type: int) -> str:
    return DTYPE_NAME.get(elem_type, f"UNKNOWN({elem_type})")


def _shape_of(value_info) -> list:
    dims = []
    for d in value_info.type.tensor_type.shape.dim:
        if d.HasField("dim_value"):
            dims.append(d.dim_value)
        elif d.HasField("dim_param") and d.dim_param:
            dims.append(d.dim_param)
        else:
            dims.append("dynamic")
    return dims


def _io_spec(value_info) -> dict:
    return {
        "name": value_info.name,
        "dtype": _elem_type_name(value_info.type.tensor_type.elem_type),
        "shape": _shape_of(value_info),
    }


def inspect(path: str) -> dict:
    size_bytes = os.path.getsize(path)
    with open(path, "rb") as fh:
        raw = fh.read()
    sha256 = hashlib.sha256(raw).hexdigest()

    model = onnx.load(path, load_external_data=False)
    graph = model.graph

    # Opset per domain.
    opset = {}
    for entry in model.opset_import:
        opset[entry.domain or "ai.onnx"] = entry.version

    # Operator histogram.
    op_hist = collections.Counter(node.op_type for node in graph.node)

    # Initializer dtype histogram + external-data detection.
    init_dtype_hist = collections.Counter()
    init_bytes = collections.Counter()
    uses_external_data = False
    for init in graph.initializer:
        dtype = _elem_type_name(init.data_type)
        init_dtype_hist[dtype] += 1
        # A tensor stored externally has data_location == EXTERNAL(1).
        if init.data_location == TensorProto.EXTERNAL:
            uses_external_data = True
        # Rough in-graph byte estimate (raw_data if present).
        init_bytes[dtype] += len(init.raw_data)

    # Sidecar external-data files sitting next to the model.
    external_files = []
    model_dir = os.path.dirname(os.path.abspath(path))
    for entry in os.listdir(model_dir) if os.path.isdir(model_dir) else []:
        if entry.endswith(".onnx_data") or entry.endswith(".data") or entry.endswith(".bin"):
            external_files.append(entry)

    # I/O contract, skipping initializers that also appear as graph inputs.
    init_names = {init.name for init in graph.initializer}
    inputs = [_io_spec(v) for v in graph.input if v.name not in init_names]
    outputs = [_io_spec(v) for v in graph.output]

    checker_ok = True
    checker_error = None
    try:
        onnx.checker.check_model(model)
    except Exception as exc:  # noqa: BLE001 - report any checker failure verbatim
        checker_ok = False
        checker_error = str(exc)

    return {
        "path": os.path.abspath(path),
        "size_bytes": size_bytes,
        "size_mib": round(size_bytes / 2**20, 3),
        "sha256": sha256,
        "ir_version": model.ir_version,
        "producer": f"{model.producer_name} {model.producer_version}".strip(),
        "opset": opset,
        "num_nodes": len(graph.node),
        "num_initializers": len(graph.initializer),
        "uses_external_data": uses_external_data,
        "external_files_in_dir": external_files,
        "initializer_dtype_histogram": dict(init_dtype_hist),
        "initializer_bytes_by_dtype": dict(init_bytes),
        "operator_histogram": dict(sorted(op_hist.items(), key=lambda kv: -kv[1])),
        "num_distinct_ops": len(op_hist),
        "inputs": inputs,
        "outputs": outputs,
        "onnx_checker_ok": checker_ok,
        "onnx_checker_error": checker_error,
    }


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("model")
    ap.add_argument("--json", dest="json_out", default=None)
    args = ap.parse_args()

    info = inspect(args.model)
    text = json.dumps(info, indent=2)
    print(text)
    if args.json_out:
        with open(args.json_out, "w") as fh:
            fh.write(text)
    return 0


if __name__ == "__main__":
    sys.exit(main())