#
#    ICRAR - International Centre for Radio Astronomy Research
#    (c) UWA - The University of Western Australia, 2016
#    Copyright by UWA (in the framework of the ICRAR)
#    All rights reserved
#
#    This library is free software; you can redistribute it and/or
#    modify it under the terms of the GNU Lesser General Public
#    License as published by the Free Software Foundation; either
#    version 2.1 of the License, or (at your option) any later version.
#
#    This library is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
#    Lesser General Public License for more details.
#
#    You should have received a copy of the GNU Lesser General Public
#    License along with this library; if not, write to the Free Software
#    Foundation, Inc., 59 Temple Place, Suite 330, Boston,
#    MA 02111-1307  USA
#
import functools
import json
import os
import tempfile
import unittest
import zlib

import six

from dlg import utils


class TestUtils(unittest.TestCase):

    def test_zlib_uncompressed_stream(self):

        fname = tempfile.mktemp()
        with open(fname, 'wb') as f:
            f.write(zlib.compress(b'abc'))

        # Read parts from the beginning
        for b,n in ((b'abc', 3), (b'ab', 2), (b'a', 1)):
            with open(fname, 'rb') as f:
                s = utils.ZlibUncompressedStream(f)
                self.assertEqual(b, s.read(n))

        # Read whole contents by successive calls
        with open(fname, 'rb') as f:
            s = utils.ZlibUncompressedStream(f)
            self.assertEqual(b'a', s.read(1))
            self.assertEqual(b'b', s.read(1))
            self.assertEqual(b'c', s.read(1))
            self.assertEqual(0, s.buflen)

        os.remove(fname)

        # Try now with bigger data sizes
        for size in [2**x+y for x in range(3,18) for y in (-1,0,1)]:

            original_bytes = os.urandom(size)
            compressed_bytes = zlib.compress(original_bytes)

            # Try first with whole reads
            compressed_stream = six.BytesIO(compressed_bytes)
            uncompressed_stream = utils.ZlibUncompressedStream(compressed_stream)
            b = uncompressed_stream.read()
            self.assertEqual(size, len(b))
            self.assertEqual(original_bytes, b)

            # Now read little by little
            read_size = min(size//4, 1024);
            b = b''
            compressed_stream = six.BytesIO(compressed_bytes)
            uncompressed_stream = utils.ZlibUncompressedStream(compressed_stream)
            for u in iter(functools.partial(uncompressed_stream.read, read_size), b''):
                b += u
            self.assertEqual(size, len(b))
            self.assertEqual(original_bytes, b)


    def test_zlib_compressed_stream_writer(self):

        compressed_ref = zlib.compress(b'abcd')

        # Read parts from the beginning
        for x in range(1, len(compressed_ref)):
            bytesio = six.BytesIO(b'abcd')
            s = utils.ZlibCompressedStream(bytesio)
            compressed = s.read(x)
            self.assertEqual(x, len(compressed))
            self.assertEqual(compressed_ref[0:x], compressed[0:x])

        # Try now with bigger data sizes
        for size in [2**x+y for x in range(3,18) for y in (-1,0,1)]:

            original_bytes = os.urandom(size)
            compressed_bytes = zlib.compress(original_bytes)

            # Try first with whole reads
            uncompressed_stream = six.BytesIO(original_bytes)
            compressed_stream = utils.ZlibCompressedStream(uncompressed_stream)
            b = compressed_stream.read()
            self.assertEqual(len(compressed_bytes), len(b), "Incorrect size when compressing %d bytes" % (size))
            self.assertEqual(compressed_bytes, b)

            # Now read little by little
            read_size = min(size//4, 1024)
            uncompressed_stream = six.BytesIO(original_bytes)
            compressed_stream = utils.ZlibCompressedStream(uncompressed_stream)
            b = b''
            for c in iter(functools.partial(compressed_stream.read, read_size), b''):
                b += c
            self.assertEqual(len(compressed_bytes), len(b), "Incorrect size when compressing %d bytes" % (size))
            self.assertEqual(compressed_bytes, b)

    def test_zlib_streams_combined_randombytes(self):
        self._test_zlib_streams_combined(os.urandom)

    def test_zlib_streams_combined_zerobytes(self):
        self._test_zlib_streams_combined(lambda n: b'0' * n)

    def _test_zlib_streams_combined(self, gen_bytes):

        sizes = [2**x+y for x in range(1,18) for y in (-1,0,1)]
        for size in sizes:

            original_bytes = gen_bytes(size)

            # Read the whole thing
            original_stream = six.BytesIO(original_bytes)
            compressed_stream = utils.ZlibCompressedStream(original_stream)
            uncompressed_stream = utils.ZlibUncompressedStream(compressed_stream)
            b = uncompressed_stream.read()
            self.assertEqual(size, len(b))
            self.assertEqual(original_bytes, b)

            # There's nothing left now to read
            for x in range(10):
                self.assertEqual(0, len(uncompressed_stream.read(100)))

            # Read with given number of bytes
            for n in (1, len(original_bytes)//2, len(original_bytes)):

                these_bytes = original_bytes[0:n]

                original_stream = six.BytesIO(these_bytes)
                compressed_stream = utils.ZlibCompressedStream(original_stream)
                uncompressed_stream = utils.ZlibUncompressedStream(compressed_stream)

                b = uncompressed_stream.read(n)
                self.assertEqual(n, len(b))
                self.assertEqual(these_bytes, b)

                # There's nothing left now to read
                for x in range(10):
                    self.assertEqual(0, len(uncompressed_stream.read(100)))

    def test_json_stream_simple_sequence(self):
        for s in ([0], [{}], ['a'], [{'oid':'A', 'type': 'plain'}]):
            stream = utils.JSONStream(s)
            self.assertEqual(s, json.loads(stream.read(100).decode('utf8')));

    def test_json_stream_sequences(self):

        ref = [1,2,3]
        objects_list = [1,2,3]
        objects_tuple = (1,2,3)
        def objects_gen():
            yield 1; yield 2; yield 3

        for objects in (objects_list, objects_tuple, objects_gen()):
            stream = utils.JSONStream(objects)
            self.assertSequenceEqual(ref, json.loads(stream.read(100).decode('latin1')))
            self.assertEqual(0, len(stream.read(100).decode('latin1')))

    def test_json_stream_simpleobj(self):

        sessionId = 'some_id'
        for obj in (1, {'a': 2}, 'b', {'sessionId': sessionId}):
            stream = utils.JSONStream(obj)
            self.assertEqual(obj, json.loads(stream.read(100).decode('latin1')))
            self.assertEqual(0, len(stream.read(100).decode('latin1')))

    def test_get_dlg_root(self):

        # It should obey the DLG_ROOT environment variable
        old = os.environ.get('DLG_ROOT', None)
        os.environ['DLG_ROOT'] = tempfile.mkdtemp()
        self.assertEqual(utils.getDlgDir(), os.environ['DLG_ROOT'])
        os.rmdir(os.environ['DLG_ROOT'])
        if old:
            os.environ['DLG_ROOT'] = old
        else:
            del os.environ['DLG_ROOT']