from __future__ import print_function, division, absolute_import

import sys

import dask
import pytest

from distributed.protocol import (loads, dumps, msgpack, maybe_compress,
                                  to_serialize)
from distributed.protocol.compression import compressions
from distributed.protocol.serialize import (Serialize, Serialized,
                                            serialize, deserialize)
from distributed.utils import nbytes
from distributed.utils_test import slow


def test_protocol():
    for msg in [1, 'a', b'a', {'x': 1}, {b'x': 1}, {'x': b''}, {}]:
        assert loads(dumps(msg)) == msg


def test_compression_1():
    pytest.importorskip('lz4')
    np = pytest.importorskip('numpy')
    x = np.ones(1000000)
    frames = dumps({'x': Serialize(x.tobytes())})
    assert sum(map(nbytes, frames)) < x.nbytes
    y = loads(frames)
    assert {'x': x.tobytes()} == y


def test_compression_2():
    pytest.importorskip('lz4')
    np = pytest.importorskip('numpy')
    x = np.random.random(10000)
    header, payload = dumps(x.tobytes())
    assert (not header or
            not msgpack.loads(header, encoding='utf8').get('compression'))


def test_compression_without_deserialization():
    pytest.importorskip('lz4')
    np = pytest.importorskip('numpy')
    x = np.ones(1000000)

    frames = dumps({'x': Serialize(x)})
    assert all(len(frame) < 1000000 for frame in frames)

    msg = loads(frames, deserialize=False)
    assert all(len(frame) < 1000000 for frame in msg['x'].frames)


def test_small():
    assert sum(map(nbytes, dumps(b''))) < 10
    assert sum(map(nbytes, dumps(1))) < 10


def test_small_and_big():
    d = {'x': (1, 2, 3), 'y': b'0' * 10000000}
    L = dumps(d)
    assert loads(L) == d
    # assert loads([small_header, small]) == {'x': [1, 2, 3]}
    # assert loads([big_header, big]) == {'y': d['y']}


def test_maybe_compress():
    pass

    try_converters = [bytes, memoryview]
    try_compressions = ['zlib', 'lz4']

    payload = b'123'

    with dask.config.set({'distributed.comm.compression': None}):
        for f in try_converters:
            assert maybe_compress(f(payload)) == (None, payload)

    for compression in try_compressions:
        try:
            __import__(compression)
        except ImportError:
            continue

        with dask.config.set({'distributed.comm.compression': compression}):
            for f in try_converters:
                payload = b'123'
                assert maybe_compress(f(payload)) == (None, payload)

                payload = b'0' * 10000
                rc, rd = maybe_compress(f(payload))
                # For some reason compressing memoryviews can force blosc...
                assert rc in (compression, 'blosc')
                assert compressions[rc]['decompress'](rd) == payload


def test_maybe_compress_sample():
    np = pytest.importorskip('numpy')
    lz4 = pytest.importorskip('lz4')
    payload = np.random.randint(0, 255, size=10000).astype('u1').tobytes()
    fmt, compressed = maybe_compress(payload)
    assert fmt is None
    assert compressed == payload


def test_large_bytes():
    for tp in (bytes, bytearray):
        msg = {'x': tp(b'0' * 1000000), 'y': 1}
        frames = dumps(msg)
        assert loads(frames) == msg
        assert len(frames[0]) < 1000
        assert len(frames[1]) < 1000

        assert loads(frames, deserialize=False) == msg


@slow
def test_large_messages():
    np = pytest.importorskip('numpy')
    psutil = pytest.importorskip('psutil')
    pytest.importorskip('lz4')
    if psutil.virtual_memory().total < 8e9:
        return

    if sys.version_info.major == 2:
        return 2

    x = np.random.randint(0, 255, size=200000000, dtype='u1')

    msg = {'x': [Serialize(x), b'small_bytes'],
           'y': {'a': Serialize(x), 'b': b'small_bytes'}}

    b = dumps(msg)
    msg2 = loads(b)
    assert msg['x'][1] == msg2['x'][1]
    assert msg['y']['b'] == msg2['y']['b']
    assert (msg['x'][0].data == msg2['x'][0]).all()
    assert (msg['y']['a'].data == msg2['y']['a']).all()


def test_large_messages_map():
    import psutil
    if psutil.virtual_memory().total < 8e9:
        pytest.skip("insufficient memory")

    x = {i: 'mystring_%d' % i for i in range(100000)}

    b = dumps(x)
    x2 = loads(b)
    assert x == x2


def test_loads_deserialize_False():
    frames = dumps({'data': Serialize(123), 'status': 'OK'})
    msg = loads(frames)
    assert msg == {'data': 123, 'status': 'OK'}

    msg = loads(frames, deserialize=False)
    assert msg['status'] == 'OK'
    assert isinstance(msg['data'], Serialized)

    result = deserialize(msg['data'].header, msg['data'].frames)
    assert result == 123


def test_loads_without_deserialization_avoids_compression():
    pytest.importorskip('lz4')
    b = b'0' * 100000

    msg = {'x': 1, 'data': to_serialize(b)}
    frames = dumps(msg)

    assert sum(map(nbytes, frames)) < 10000

    msg2 = loads(frames, deserialize=False)
    assert sum(map(nbytes, msg2['data'].frames)) < 10000

    msg3 = dumps(msg2)
    msg4 = loads(msg3)

    assert msg4 == {'x': 1, 'data': b'0' * 100000}


def eq_frames(a, b):
    if b'headers' in a:
        return (msgpack.loads(a, use_list=False)
                == msgpack.loads(b, use_list=False))
    else:
        return a == b


def test_dumps_loads_Serialize():
    msg = {'x': 1, 'data': Serialize(123)}
    frames = dumps(msg)
    assert len(frames) > 2
    result = loads(frames)
    assert result == {'x': 1, 'data': 123}

    result2 = loads(frames, deserialize=False)
    assert result2['x'] == 1
    assert isinstance(result2['data'], Serialized)
    assert any(a is b
               for a in result2['data'].frames
               for b in frames)

    frames2 = dumps(result2)
    assert all(map(eq_frames, frames, frames2))

    result3 = loads(frames2)
    assert result == result3


def test_dumps_loads_Serialized():
    msg = {'x': 1,
           'data': Serialized(*serialize(123)),
           }
    frames = dumps(msg)
    assert len(frames) > 2
    result = loads(frames)
    assert result == {'x': 1, 'data': 123}

    result2 = loads(frames, deserialize=False)
    assert result2 == msg

    frames2 = dumps(result2)
    assert all(map(eq_frames, frames, frames2))

    result3 = loads(frames2)
    assert result == result3


@pytest.mark.skipif(sys.version_info[0] < 3,
                    reason='NumPy doesnt use memoryviews')
def test_maybe_compress_memoryviews():
    np = pytest.importorskip('numpy')
    pytest.importorskip('lz4')
    x = np.arange(1000000, dtype='int64')
    compression, payload = maybe_compress(x.data)
    try:
        import blosc  # noqa: F401
    except ImportError:
        assert compression == 'lz4'
        assert len(payload) < x.nbytes * 0.75
    else:
        assert compression == 'blosc'
        assert len(payload) < x.nbytes / 10