import math
import struct

import msgpack
import numpy as np

from . import sharray


def pack(data):
    leaves, treedef = tree_flatten(data)
    specs, buffers = [], []
    for value in leaves:
        if value is None:
            specs.append(['none'])
            buffers.append(b'\x00')
        elif isinstance(value, str):
            specs.append(['utf8'])
            buffers.append(value.encode('utf-8'))
        elif isinstance(value, (bytes, bytearray, memoryview)):
            if isinstance(value, memoryview):
                value = value.cast('c')
                assert value.c_contiguous
            specs.append(['bytes'])
            buffers.append(value)
        elif isinstance(value, (np.ndarray, np.generic, int, float)):
            value = np.asarray(value)
            if value.dtype == object:
                raise TypeError(data)
            assert value.data.c_contiguous, (
                'Array is not contiguous in memory. Use '
                + "np.asarray(arr, order='C') before passing the data into pack()."
            )
            specs.append(['array', value.shape, value.dtype.str])
            buffer = value.data.cast('c') if value.size else b'\x00'
            buffers.append(buffer)
        elif isinstance(value, sharray.SharedArray):
            specs.append(['sharray', *value.__getstate__()])
            buffers.append(b'\x00')
        else:
            raise NotImplementedError(type(value))
    buffers = [msgpack.packb(treedef), msgpack.packb(specs), *buffers]
    length = len(buffers).to_bytes(8, 'little', signed=False)
    sizes = struct.pack('<' + ('Q' * len(buffers)), *[len(x) for x in buffers])
    buffers = [length, sizes, *buffers]
    return buffers


def unpack(buffer):
    length = int.from_bytes(buffer[:8], 'little', signed=False)
    buffer = buffer[8:]
    sizes = struct.unpack('<' + ('Q' * length), buffer[: 8 * length])
    buffer = buffer[8 * length :]
    limits = np.cumsum(sizes)
    buffers = [buffer[i:j] for i, j in zip([0, *limits[:-1]], limits)]
    treedef, specs, *buffers = buffers
    treedef = msgpack.unpackb(treedef)
    specs = msgpack.unpackb(specs)
    leaves = []
    for spec, buffer in zip(specs, buffers):
        if spec[0] == 'none':
            assert buffer == b'\x00'
            leaves.append(None)
        elif spec[0] == 'utf8':
            leaves.append(bytes(buffer).decode('utf-8'))
        elif spec[0] == 'bytes':
            leaves.append(buffer)
        elif spec[0] == 'array':
            shape, dtype = spec[1:]
            if not math.prod(shape):
                assert buffer == b'\x00'
                buffer = b''
            leaves.append(np.frombuffer(buffer, dtype).reshape(shape))
        elif spec[0] == 'sharray':
            assert buffer == b'\x00'
            leaves.append(sharray.SharedArray(*spec[1:]))
        else:
            raise NotImplementedError(spec)
    data = tree_unflatten(leaves, treedef)
    return data


def tree_map(fn, *trees, isleaf=None):
    assert trees, 'Provide one or more nested Python structures'
    kw = dict(isleaf=isleaf)
    first = trees[0]
    try:
        assert all(isinstance(x, type(first)) for x in trees)
        if isleaf and isleaf(trees[0]):
            return fn(*trees)
        if isinstance(first, list):
            assert all(len(x) == len(first) for x in trees)
            return [
                tree_map(fn, *[t[i] for t in trees], **kw)
                for i in range(len(first))
            ]
        if isinstance(first, tuple):
            assert all(len(x) == len(first) for x in trees)
            return tuple(
                tree_map(fn, *[t[i] for t in trees], **kw)
                for i in range(len(first))
            )
        if isinstance(first, dict):
            assert all(set(x.keys()) == set(first.keys()) for x in trees)
            return {
                k: tree_map(fn, *[t[k] for t in trees], **kw) for k in first
            }
        if hasattr(first, 'keys') and hasattr(first, 'get'):
            assert all(set(x.keys()) == set(first.keys()) for x in trees)
            return type(first)(
                {k: tree_map(fn, *[t[k] for t in trees], **kw) for k in first}
            )
    except AssertionError:
        raise TypeError('Tree structure differs between arguments')
    return fn(*trees)


def tree_flatten(tree, isleaf=None):
    leaves = []
    tree_map(lambda x: leaves.append(x), tree, isleaf=isleaf)
    structure = tree_map(lambda x: None, tree, isleaf=isleaf)
    return tuple(leaves), structure


def tree_unflatten(leaves, structure):
    leaves = iter(tuple(leaves))
    return tree_map(lambda x: next(leaves), structure)


def tree_equals(xs, ys):
    assert type(xs) == type(ys), (type(xs), type(ys))
    if isinstance(xs, (list, tuple)):
        assert len(xs) == len(ys)
        return all(tree_equals(x, y) for x, y in zip(xs, ys))
    elif isinstance(xs, dict):
        assert xs.keys() == ys.keys()
        return all(tree_equals(xs[k], ys[k]) for k in xs.keys())
    elif isinstance(xs, np.ndarray):
        assert xs.shape == ys.shape, (xs.shape, ys.shape)
        assert xs.dtype == ys.dtype, (xs.dtype, ys.dtype)
        return (xs == ys).all()
    elif isinstance(xs, sharray.SharedArray):
        assert xs.name == ys.name, (xs.name, ys.name)
        info = (xs.array.shape, ys.array.shape)
        assert xs.array.shape == ys.array.shape, info
        assert xs.array.dtype == ys.array.dtype, info
        return (xs == ys).all()
    else:
        return xs == ys
