"""
onnx_to_art.py  —  AllayRT model converter
============================================
Converts a .onnx model to the AllayRT .art binary format.

Usage (CLI):
    python onnx_to_art.py model.onnx model.art

Usage (Python):
    from onnx_to_art import convert
    convert("model.onnx", "model.art")
"""

import argparse
import numpy as np

from AllayRT.model import ARTModel, save_art, ONNX_DTYPE_MAP


# ─────────────────────────────────────────────
#  Attribute helpers
# ─────────────────────────────────────────────

def _attr_value(attr):
    """Extract a Python value from an ONNX AttributeProto."""
    from onnx import AttributeProto  # type: ignore
    t = attr.type
    if t == AttributeProto.FLOAT:      return float(attr.f)
    if t == AttributeProto.INT:        return int(attr.i)
    if t == AttributeProto.STRING:     return attr.s.decode("utf-8")
    if t == AttributeProto.TENSOR:     return _tensor_to_numpy(attr.t).tolist()
    if t == AttributeProto.FLOATS:     return list(attr.floats)
    if t == AttributeProto.INTS:       return list(attr.ints)
    if t == AttributeProto.STRINGS:    return [s.decode("utf-8") for s in attr.strings]
    if t == AttributeProto.GRAPH:      return "<subgraph>"
    return None


def _tensor_to_numpy(tensor) -> np.ndarray:
    """Convert an ONNX TensorProto to numpy array."""
    import onnx.numpy_helper as nph  # type: ignore
    return nph.to_array(tensor)


def _type_proto_info(type_proto) -> dict:
    """Extract shape/dtype info from a TypeProto."""
    info = {"dtype": "float32", "shape": None}
    if type_proto.HasField("tensor_type"):
        tt = type_proto.tensor_type
        elem = tt.elem_type
        np_dt = ONNX_DTYPE_MAP.get(elem, np.float32)
        info["dtype"] = np.dtype(np_dt).str.lstrip("<>=!")
        if tt.HasField("shape"):
            dims = []
            for d in tt.shape.dim:
                if d.HasField("dim_value"):
                    dims.append(d.dim_value)
                elif d.HasField("dim_param"):
                    dims.append(d.dim_param)  # symbolic e.g. "batch_size"
                else:
                    dims.append(None)
            info["shape"] = dims
    return info


# ─────────────────────────────────────────────
#  Main converter
# ─────────────────────────────────────────────

def convert(onnx_path: str, art_path: str, verbose: bool = True) -> ARTModel:
    """
    Load an ONNX model and write it as a .art file.

    Parameters
    ----------
    onnx_path : path to source .onnx file
    art_path  : path for output .art file
    verbose   : print progress info

    Returns
    -------
    ARTModel  : the converted in-memory model
    """
    try:
        import onnx  # type: ignore
    except ImportError:
        raise ImportError(
            "[AllayRT] 'onnx' package not found. Install it with:\n"
            "    pip install onnx"
        )

    if verbose:
        print(f"[AllayRT Converter] Reading {onnx_path} ...")

    onnx_model = onnx.load(onnx_path)
    onnx.checker.check_model(onnx_model)

    graph = onnx_model.graph
    art = ARTModel()

    # ── Metadata ──────────────────────────────
    opset_ver = (
        onnx_model.opset_import[0].version
        if onnx_model.opset_import else 0
    )
    art.meta = {
        "name":           onnx_model.graph.name or "unnamed",
        "ir_version":     onnx_model.ir_version,
        "opset_version":  opset_ver,
        "producer_name":  onnx_model.producer_name,
        "producer_version": onnx_model.producer_version,
        "domain":         onnx_model.domain,
        "model_version":  onnx_model.model_version,
        "doc_string":     onnx_model.doc_string,
        "converted_by":   "AllayRT onnx_to_art v1.0",
    }

    # ── Initializers (weights) ─────────────────
    init_names = set()
    for init in graph.initializer:
        arr = _tensor_to_numpy(init)
        art.tensors[init.name] = arr
        init_names.add(init.name)
        if verbose:
            print(f"  tensor  {init.name:40s}  {str(arr.shape):20s}  {arr.dtype}")

    # ── Graph inputs ──────────────────────────
    for inp in graph.input:
        if inp.name in init_names:
            continue  # initializers appear as inputs in older ONNX models
        info = _type_proto_info(inp.type)
        art.inputs.append({"name": inp.name, **info})

    # ── Graph outputs ─────────────────────────
    for out in graph.output:
        info = _type_proto_info(out.type)
        art.outputs.append({"name": out.name, **info})

    # ── Value info (intermediate shapes) ──────
    for vi in graph.value_info:
        info = _type_proto_info(vi.type)
        art.value_info[vi.name] = info

    # ── Nodes ─────────────────────────────────
    for n in graph.node:
        node_dict = {
            "op_type": n.op_type,
            "name":    n.name,
            "domain":  n.domain,
            "inputs":  list(n.input),
            "outputs": list(n.output),
            "attrs":   {a.name: _attr_value(a) for a in n.attribute},
        }
        art.nodes.append(node_dict)
        if verbose:
            ins  = ", ".join(n.input)
            outs = ", ".join(n.output)
            print(f"  node    {n.op_type:20s}  [{ins}]  ->  [{outs}]")

    # ── Write .art ────────────────────────────
    save_art(art, art_path)

    if verbose:
        print(f"\n[AllayRT Converter] ✓ Conversion complete.")
        print(f"  Source : {onnx_path}")
        print(f"  Output : {art_path}")
        print(f"  Model  : {art}")

    return art


# ─────────────────────────────────────────────
#  CLI entry point
# ─────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(
        description="AllayRT: Convert ONNX model → .art format"
    )
    parser.add_argument("onnx_path",  help="Input .onnx model path")
    parser.add_argument("art_path",   help="Output .art model path")
    parser.add_argument("--quiet", "-q", action="store_true",
                        help="Suppress verbose output")
    args = parser.parse_args()

    convert(args.onnx_path, args.art_path, verbose=not args.quiet)


if __name__ == "__main__":
    main()
