import struct
import numpy as np

# --- The Tenso Protocol ---
_MAGIC = b'TNSO'  # Magic bytes for file identification
_VERSION = 1

_DTYPE_MAP = {
    np.dtype('float32'): 1,
    np.dtype('int32'): 2,
    np.dtype('float64'): 3,
    np.dtype('int64'): 4,
}
_REV_DTYPE_MAP = {v: k for k, v in _DTYPE_MAP.items()}

def dumps(tensor: np.ndarray) -> bytes:
    """Serialize a numpy array into bytes (Zero-Copy)."""
    # 1. Validation
    if tensor.dtype not in _DTYPE_MAP:
        raise ValueError(f"Unsupported dtype: {tensor.dtype}")
    
    dtype_code = _DTYPE_MAP[tensor.dtype]
    shape = tensor.shape
    ndim = len(shape)
    
    # 2. Header (8 Bytes): Magic + Ver + Flags + Dtype + Ndim
    header = struct.pack('<4sBBBB', _MAGIC, _VERSION, 0, dtype_code, ndim)
    
    # 3. Shape Block (Variable): Ndim * uint32
    shape_block = struct.pack(f'<{ndim}I', *shape)
    
    # 4. Body (Zero-Copy): Raw memory dump
    return header + shape_block + tensor.tobytes()

def loads(data: bytes) -> np.ndarray:
    """Deserialize bytes back into a numpy array."""
    # 1. Parse Header
    magic, ver, flags, dtype_code, ndim = struct.unpack('<4sBBBB', data[:8])
    
    if magic != _MAGIC:
        raise ValueError("Invalid tenso packet (Magic bytes mismatch)")
    
    # 2. Parse Shape
    shape_end = 8 + (ndim * 4)
    shape = struct.unpack(f'<{ndim}I', data[8:shape_end])
    
    # 3. Parse Body
    dtype = _REV_DTYPE_MAP[dtype_code]
    return np.frombuffer(data, dtype=dtype, offset=shape_end).reshape(shape)

# File I/O Helpers
def dump(tensor: np.ndarray, fp) -> None:
    fp.write(dumps(tensor))

def load(fp) -> np.ndarray:
    return loads(fp.read())