"""
AllayRT Model Format (.art)
===========================

Binary layout:
  [MAGIC 8B] [VERSION 4B] [FLAGS 4B] [HEADER_LEN 4B] [HEADER JSON]
  [N_TENSORS 4B]
  For each tensor:
    [NAME_LEN 4B] [NAME bytes]
    [DTYPE 1B] [NDIM 1B] [SHAPE: ndim * 8B]
    [DATA_LEN 8B] [DATA bytes]
  [N_NODES 4B]
  For each node:
    [JSON_LEN 4B] [JSON bytes]
  [CHECKSUM 4B]  -- CRC32 of everything above

DTYPE encoding:
  0x01 = float32
  0x02 = float64
  0x03 = int32
  0x04 = int64
  0x05 = int8
  0x06 = uint8
  0x07 = bool
  0x08 = float16
"""

import struct
import json
import zlib
import numpy as np

MAGIC = b"ALLAYRT\x00"
VERSION = (1, 0, 0, 0)  # major, minor, patch, reserved
FORMAT_VERSION_INT = 0x00010000

DTYPE_MAP = {
    np.dtype("float32"): 0x01,
    np.dtype("float64"): 0x02,
    np.dtype("int32"):   0x03,
    np.dtype("int64"):   0x04,
    np.dtype("int8"):    0x05,
    np.dtype("uint8"):   0x06,
    np.dtype("bool"):    0x07,
    np.dtype("float16"): 0x08,
}

DTYPE_MAP_INV = {v: k for k, v in DTYPE_MAP.items()}

# ONNX elem_type -> numpy dtype
ONNX_DTYPE_MAP = {
    1:  np.float32,
    2:  np.uint8,
    3:  np.int8,
    4:  np.uint16,
    5:  np.int16,
    6:  np.int32,
    7:  np.int64,
    9:  bool,
    10: np.float16,
    11: np.float64,
    12: np.uint32,
    13: np.uint64,
}


class ARTModel:
    """In-memory representation of a .art model."""

    def __init__(self):
        self.meta: dict = {}          # model metadata (name, opset, etc.)
        self.tensors: dict = {}       # name -> np.ndarray  (weights/initializers)
        self.nodes: list = []         # list of op dicts
        self.inputs: list = []        # list of {name, shape, dtype}
        self.outputs: list = []       # list of {name, shape, dtype}
        self.value_info: dict = {}    # intermediate tensor shapes

    def __repr__(self):
        return (
            f"<ARTModel name={self.meta.get('name','?')} "
            f"nodes={len(self.nodes)} tensors={len(self.tensors)} "
            f"inputs={[i['name'] for i in self.inputs]} "
            f"outputs={[o['name'] for o in self.outputs]}>"
        )


# ─────────────────────────────────────────────
#  Writer
# ─────────────────────────────────────────────

def save_art(model: ARTModel, path: str):
    """Serialize an ARTModel to a .art file."""
    buf = bytearray()

    # MAGIC + VERSION + FLAGS
    buf += MAGIC
    buf += struct.pack("<I", FORMAT_VERSION_INT)
    buf += struct.pack("<I", 0)  # flags reserved

    # HEADER JSON
    header = {
        "meta": model.meta,
        "inputs": model.inputs,
        "outputs": model.outputs,
        "value_info": model.value_info,
    }
    header_bytes = json.dumps(header, separators=(",", ":")).encode("utf-8")
    buf += struct.pack("<I", len(header_bytes))
    buf += header_bytes

    # TENSORS
    buf += struct.pack("<I", len(model.tensors))
    for name, arr in model.tensors.items():
        arr = np.ascontiguousarray(arr)
        name_b = name.encode("utf-8")
        buf += struct.pack("<I", len(name_b))
        buf += name_b
        dtype_code = DTYPE_MAP.get(arr.dtype, 0x01)
        buf += struct.pack("<B", dtype_code)
        buf += struct.pack("<B", arr.ndim)
        for dim in arr.shape:
            buf += struct.pack("<q", dim)
        raw = arr.tobytes()
        buf += struct.pack("<Q", len(raw))
        buf += raw

    # NODES
    buf += struct.pack("<I", len(model.nodes))
    for node in model.nodes:
        node_b = json.dumps(node, separators=(",", ":")).encode("utf-8")
        buf += struct.pack("<I", len(node_b))
        buf += node_b

    # CHECKSUM
    crc = zlib.crc32(buf) & 0xFFFFFFFF
    buf += struct.pack("<I", crc)

    with open(path, "wb") as f:
        f.write(buf)

    size_kb = len(buf) / 1024
    print(f"[AllayRT] Saved {path}  ({size_kb:.1f} KB, "
          f"{len(model.tensors)} tensors, {len(model.nodes)} nodes)")


# ─────────────────────────────────────────────
#  Reader
# ─────────────────────────────────────────────

def load_art(path: str) -> ARTModel:
    """Deserialize a .art file into an ARTModel."""
    with open(path, "rb") as f:
        data = f.read()

    pos = 0

    # Verify checksum
    stored_crc = struct.unpack_from("<I", data, len(data) - 4)[0]
    computed_crc = zlib.crc32(data[:-4]) & 0xFFFFFFFF
    if stored_crc != computed_crc:
        raise ValueError(
            f"[AllayRT] Checksum mismatch in {path}: "
            f"expected {computed_crc:#010x}, got {stored_crc:#010x}"
        )

    # MAGIC
    magic = data[pos:pos+8]; pos += 8
    if magic != MAGIC:
        raise ValueError(f"[AllayRT] Not a valid .art file (bad magic): {magic}")

    # VERSION
    ver = struct.unpack_from("<I", data, pos)[0]; pos += 4
    major = (ver >> 16) & 0xFF
    minor = (ver >> 8) & 0xFF

    # FLAGS
    _flags = struct.unpack_from("<I", data, pos)[0]; pos += 4

    # HEADER JSON
    hdr_len = struct.unpack_from("<I", data, pos)[0]; pos += 4
    header = json.loads(data[pos:pos+hdr_len].decode("utf-8")); pos += hdr_len

    model = ARTModel()
    model.meta = header.get("meta", {})
    model.inputs = header.get("inputs", [])
    model.outputs = header.get("outputs", [])
    model.value_info = header.get("value_info", {})

    # TENSORS
    n_tensors = struct.unpack_from("<I", data, pos)[0]; pos += 4
    for _ in range(n_tensors):
        nlen = struct.unpack_from("<I", data, pos)[0]; pos += 4
        name = data[pos:pos+nlen].decode("utf-8"); pos += nlen
        dtype_code = struct.unpack_from("<B", data, pos)[0]; pos += 1
        ndim = struct.unpack_from("<B", data, pos)[0]; pos += 1
        shape = tuple(struct.unpack_from("<q", data, pos + i*8)[0] for i in range(ndim))
        pos += ndim * 8
        dlen = struct.unpack_from("<Q", data, pos)[0]; pos += 8
        raw = data[pos:pos+dlen]; pos += dlen
        dtype = DTYPE_MAP_INV.get(dtype_code, np.dtype("float32"))
        arr = np.frombuffer(raw, dtype=dtype).reshape(shape).copy()
        model.tensors[name] = arr

    # NODES
    n_nodes = struct.unpack_from("<I", data, pos)[0]; pos += 4
    for _ in range(n_nodes):
        jlen = struct.unpack_from("<I", data, pos)[0]; pos += 4
        node = json.loads(data[pos:pos+jlen].decode("utf-8")); pos += jlen
        model.nodes.append(node)

    print(f"[AllayRT] Loaded {path}  "
          f"(v{major}.{minor}, {len(model.tensors)} tensors, {len(model.nodes)} nodes)")
    return model
