import io

import os

import random

import struct

import sys

import tempfile

import unittest

import zstandard as zstd

from .common import (

    generate_samples,

    make_cffi,

    NonClosingBytesIO,

    OpCountingBytesIO,

)

if sys.version_info[0] >= 3:

    next = lambda it: it.__next__()

else:

    next = lambda it: it.next()

@make_cffi

class TestFrameHeaderSize(unittest.TestCase):

    def test_empty(self):

        with self.assertRaisesRegexp(

            zstd.ZstdError, 'could not determine frame header size: Src size '

                            'is incorrect'):

            zstd.frame_header_size(b'')

    def test_too_small(self):

        with self.assertRaisesRegexp(

            zstd.ZstdError, 'could not determine frame header size: Src size '

                            'is incorrect'):

            zstd.frame_header_size(b'foob')

    def test_basic(self):

        # It doesn't matter that it isn't a valid frame.

        self.assertEqual(zstd.frame_header_size(b'long enough but no magic'), 6)

@make_cffi

class TestFrameContentSize(unittest.TestCase):

    def test_empty(self):

        with self.assertRaisesRegexp(zstd.ZstdError,

                                     'error when determining content size'):

            zstd.frame_content_size(b'')

    def test_too_small(self):

        with self.assertRaisesRegexp(zstd.ZstdError,

                                     'error when determining content size'):

            zstd.frame_content_size(b'foob')

    def test_bad_frame(self):

        with self.assertRaisesRegexp(zstd.ZstdError,

                                     'error when determining content size'):

            zstd.frame_content_size(b'invalid frame header')

    def test_unknown(self):

        cctx = zstd.ZstdCompressor(write_content_size=False)

        frame = cctx.compress(b'foobar')

        self.assertEqual(zstd.frame_content_size(frame), -1)

    def test_empty(self):

        cctx = zstd.ZstdCompressor()

        frame = cctx.compress(b'')

        self.assertEqual(zstd.frame_content_size(frame), 0)

    def test_basic(self):

        cctx = zstd.ZstdCompressor()

        frame = cctx.compress(b'foobar')

        self.assertEqual(zstd.frame_content_size(frame), 6)

@make_cffi

class TestDecompressor(unittest.TestCase):

    def test_memory_size(self):

        dctx = zstd.ZstdDecompressor()

        self.assertGreater(dctx.memory_size(), 100)

@make_cffi

class TestDecompressor_decompress(unittest.TestCase):

    def test_empty_input(self):

        dctx = zstd.ZstdDecompressor()

        with self.assertRaisesRegexp(zstd.ZstdError, 'error determining content size from frame header'):

            dctx.decompress(b'')

    def test_invalid_input(self):

        dctx = zstd.ZstdDecompressor()

        with self.assertRaisesRegexp(zstd.ZstdError, 'error determining content size from frame header'):

            dctx.decompress(b'foobar')

    def test_input_types(self):

        cctx = zstd.ZstdCompressor(level=1)

        compressed = cctx.compress(b'foo')

        mutable_array = bytearray(len(compressed))

        mutable_array[:] = compressed

        sources = [

            memoryview(compressed),

            bytearray(compressed),

            mutable_array,

        ]

        dctx = zstd.ZstdDecompressor()

        for source in sources:

            self.assertEqual(dctx.decompress(source), b'foo')

    def test_no_content_size_in_frame(self):

        cctx = zstd.ZstdCompressor(write_content_size=False)

        compressed = cctx.compress(b'foobar')

        dctx = zstd.ZstdDecompressor()

        with self.assertRaisesRegexp(zstd.ZstdError, 'could not determine content size in frame header'):

            dctx.decompress(compressed)

    def test_content_size_present(self):

        cctx = zstd.ZstdCompressor()

        compressed = cctx.compress(b'foobar')

        dctx = zstd.ZstdDecompressor()

        decompressed = dctx.decompress(compressed)

        self.assertEqual(decompressed, b'foobar')

    def test_empty_roundtrip(self):

        cctx = zstd.ZstdCompressor()

        compressed = cctx.compress(b'')

        dctx = zstd.ZstdDecompressor()

        decompressed = dctx.decompress(compressed)

        self.assertEqual(decompressed, b'')

    def test_max_output_size(self):

        cctx = zstd.ZstdCompressor(write_content_size=False)

        source = b'foobar' * 256

        compressed = cctx.compress(source)

        dctx = zstd.ZstdDecompressor()

        # Will fit into buffer exactly the size of input.

        decompressed = dctx.decompress(compressed, max_output_size=len(source))

        self.assertEqual(decompressed, source)

        # Input size - 1 fails

        with self.assertRaisesRegexp(zstd.ZstdError,

                'decompression error: did not decompress full frame'):

            dctx.decompress(compressed, max_output_size=len(source) - 1)

        # Input size + 1 works

        decompressed = dctx.decompress(compressed, max_output_size=len(source) + 1)

        self.assertEqual(decompressed, source)

        # A much larger buffer works.

        decompressed = dctx.decompress(compressed, max_output_size=len(source) * 64)

        self.assertEqual(decompressed, source)

    def test_stupidly_large_output_buffer(self):

        cctx = zstd.ZstdCompressor(write_content_size=False)

        compressed = cctx.compress(b'foobar' * 256)

        dctx = zstd.ZstdDecompressor()

        # Will get OverflowError on some Python distributions that can't

        # handle really large integers.

        with self.assertRaises((MemoryError, OverflowError)):

            dctx.decompress(compressed, max_output_size=2**62)

    def test_dictionary(self):

        samples = []

        for i in range(128):

            samples.append(b'foo' * 64)

            samples.append(b'bar' * 64)

            samples.append(b'foobar' * 64)

        d = zstd.train_dictionary(8192, samples)

        orig = b'foobar' * 16384

        cctx = zstd.ZstdCompressor(level=1, dict_data=d)

        compressed = cctx.compress(orig)

        dctx = zstd.ZstdDecompressor(dict_data=d)

        decompressed = dctx.decompress(compressed)

        self.assertEqual(decompressed, orig)

    def test_dictionary_multiple(self):

        samples = []

        for i in range(128):

            samples.append(b'foo' * 64)

            samples.append(b'bar' * 64)

            samples.append(b'foobar' * 64)

        d = zstd.train_dictionary(8192, samples)

        sources = (b'foobar' * 8192, b'foo' * 8192, b'bar' * 8192)

        compressed = []

        cctx = zstd.ZstdCompressor(level=1, dict_data=d)

        for source in sources:

            compressed.append(cctx.compress(source))

        dctx = zstd.ZstdDecompressor(dict_data=d)

        for i in range(len(sources)):

            decompressed = dctx.decompress(compressed[i])

            self.assertEqual(decompressed, sources[i])

    def test_max_window_size(self):

        with open(__file__, 'rb') as fh:

            source = fh.read()

        # If we write a content size, the decompressor engages single pass

        # mode and the window size doesn't come into play.

        cctx = zstd.ZstdCompressor(write_content_size=False)

        frame = cctx.compress(source)

        dctx = zstd.ZstdDecompressor(max_window_size=2**zstd.WINDOWLOG_MIN)

        with self.assertRaisesRegexp(

            zstd.ZstdError, 'decompression error: Frame requires too much memory'):

            dctx.decompress(frame, max_output_size=len(source))

@make_cffi

class TestDecompressor_copy_stream(unittest.TestCase):

    def test_no_read(self):

        source = object()

        dest = io.BytesIO()

        dctx = zstd.ZstdDecompressor()

        with self.assertRaises(ValueError):

            dctx.copy_stream(source, dest)

    def test_no_write(self):

        source = io.BytesIO()

        dest = object()

        dctx = zstd.ZstdDecompressor()

        with self.assertRaises(ValueError):

            dctx.copy_stream(source, dest)

    def test_empty(self):

        source = io.BytesIO()

        dest = io.BytesIO()

        dctx = zstd.ZstdDecompressor()

        # TODO should this raise an error?

        r, w = dctx.copy_stream(source, dest)

        self.assertEqual(r, 0)

        self.assertEqual(w, 0)

        self.assertEqual(dest.getvalue(), b'')

    def test_large_data(self):

        source = io.BytesIO()

        for i in range(255):

            source.write(struct.Struct('>B').pack(i) * 16384)

        source.seek(0)

        compressed = io.BytesIO()

        cctx = zstd.ZstdCompressor()

        cctx.copy_stream(source, compressed)

        compressed.seek(0)

        dest = io.BytesIO()

        dctx = zstd.ZstdDecompressor()

        r, w = dctx.copy_stream(compressed, dest)

        self.assertEqual(r, len(compressed.getvalue()))

        self.assertEqual(w, len(source.getvalue()))

    def test_read_write_size(self):

        source = OpCountingBytesIO(zstd.ZstdCompressor().compress(

            b'foobarfoobar'))

        dest = OpCountingBytesIO()

        dctx = zstd.ZstdDecompressor()

        r, w = dctx.copy_stream(source, dest, read_size=1, write_size=1)

        self.assertEqual(r, len(source.getvalue()))

        self.assertEqual(w, len(b'foobarfoobar'))

        self.assertEqual(source._read_count, len(source.getvalue()) + 1)

        self.assertEqual(dest._write_count, len(dest.getvalue()))

@make_cffi

class TestDecompressor_stream_reader(unittest.TestCase):

    def test_context_manager(self):

        dctx = zstd.ZstdDecompressor()

        with dctx.stream_reader(b'foo') as reader:

            with self.assertRaisesRegexp(ValueError, 'cannot __enter__ multiple times'):

                with reader as reader2:

                    pass

    def test_not_implemented(self):

        dctx = zstd.ZstdDecompressor()

        with dctx.stream_reader(b'foo') as reader:

            with self.assertRaises(io.UnsupportedOperation):

                reader.readline()

            with self.assertRaises(io.UnsupportedOperation):

                reader.readlines()

            with self.assertRaises(io.UnsupportedOperation):

                iter(reader)

            with self.assertRaises(io.UnsupportedOperation):

                next(reader)

            with self.assertRaises(io.UnsupportedOperation):

                reader.write(b'foo')

            with self.assertRaises(io.UnsupportedOperation):

                reader.writelines([])

    def test_constant_methods(self):

        dctx = zstd.ZstdDecompressor()

        with dctx.stream_reader(b'foo') as reader:

            self.assertFalse(reader.closed)

            self.assertTrue(reader.readable())

            self.assertFalse(reader.writable())

            self.assertTrue(reader.seekable())

            self.assertFalse(reader.isatty())

            self.assertFalse(reader.closed)

            self.assertIsNone(reader.flush())

            self.assertFalse(reader.closed)

        self.assertTrue(reader.closed)

    def test_read_closed(self):

        dctx = zstd.ZstdDecompressor()

        with dctx.stream_reader(b'foo') as reader:

            reader.close()

            self.assertTrue(reader.closed)

            with self.assertRaisesRegexp(ValueError, 'stream is closed'):

                reader.read(1)

    def test_read_sizes(self):

        cctx = zstd.ZstdCompressor()

        foo = cctx.compress(b'foo')

        dctx = zstd.ZstdDecompressor()

        with dctx.stream_reader(foo) as reader:

            with self.assertRaisesRegexp(ValueError, 'cannot read negative amounts less than -1'):

                reader.read(-2)

            self.assertEqual(reader.read(0), b'')

            self.assertEqual(reader.read(), b'foo')

    def test_read_buffer(self):

        cctx = zstd.ZstdCompressor()

        source = b''.join([b'foo' * 60, b'bar' * 60, b'baz' * 60])

        frame = cctx.compress(source)

        dctx = zstd.ZstdDecompressor()

        with dctx.stream_reader(frame) as reader:

            self.assertEqual(reader.tell(), 0)

            # We should get entire frame in one read.

            result = reader.read(8192)

            self.assertEqual(result, source)

            self.assertEqual(reader.tell(), len(source))

            # Read after EOF should return empty bytes.

            self.assertEqual(reader.read(1), b'')

            self.assertEqual(reader.tell(), len(result))

        self.assertTrue(reader.closed)

    def test_read_buffer_small_chunks(self):

        cctx = zstd.ZstdCompressor()

        source = b''.join([b'foo' * 60, b'bar' * 60, b'baz' * 60])

        frame = cctx.compress(source)

        dctx = zstd.ZstdDecompressor()

        chunks = []

        with dctx.stream_reader(frame, read_size=1) as reader:

            while True:

                chunk = reader.read(1)

                if not chunk:

                    break

                chunks.append(chunk)

                self.assertEqual(reader.tell(), sum(map(len, chunks)))

        self.assertEqual(b''.join(chunks), source)

    def test_read_stream(self):

        cctx = zstd.ZstdCompressor()

        source = b''.join([b'foo' * 60, b'bar' * 60, b'baz' * 60])

        frame = cctx.compress(source)

        dctx = zstd.ZstdDecompressor()

        with dctx.stream_reader(io.BytesIO(frame)) as reader:

            self.assertEqual(reader.tell(), 0)

            chunk = reader.read(8192)

            self.assertEqual(chunk, source)

            self.assertEqual(reader.tell(), len(source))

            self.assertEqual(reader.read(1), b'')

            self.assertEqual(reader.tell(), len(source))

            self.assertFalse(reader.closed)

        self.assertTrue(reader.closed)

    def test_read_stream_small_chunks(self):

        cctx = zstd.ZstdCompressor()

        source = b''.join([b'foo' * 60, b'bar' * 60, b'baz' * 60])

        frame = cctx.compress(source)

        dctx = zstd.ZstdDecompressor()

        chunks = []

        with dctx.stream_reader(io.BytesIO(frame), read_size=1) as reader:

            while True:

                chunk = reader.read(1)

                if not chunk:

                    break

                chunks.append(chunk)

                self.assertEqual(reader.tell(), sum(map(len, chunks)))

        self.assertEqual(b''.join(chunks), source)

    def test_read_after_exit(self):

        cctx = zstd.ZstdCompressor()

        frame = cctx.compress(b'foo' * 60)

        dctx = zstd.ZstdDecompressor()

        with dctx.stream_reader(frame) as reader:

            while reader.read(16):

                pass

        self.assertTrue(reader.closed)

        with self.assertRaisesRegexp(ValueError, 'stream is closed'):

            reader.read(10)

    def test_illegal_seeks(self):

        cctx = zstd.ZstdCompressor()

        frame = cctx.compress(b'foo' * 60)

        dctx = zstd.ZstdDecompressor()

        with dctx.stream_reader(frame) as reader:

            with self.assertRaisesRegexp(ValueError,

                                         'cannot seek to negative position'):

                reader.seek(-1, os.SEEK_SET)

            reader.read(1)

            with self.assertRaisesRegexp(

                ValueError, 'cannot seek zstd decompression stream backwards'):

                reader.seek(0, os.SEEK_SET)

            with self.assertRaisesRegexp(

                ValueError, 'cannot seek zstd decompression stream backwards'):

                reader.seek(-1, os.SEEK_CUR)

            with self.assertRaisesRegexp(

                ValueError,

                'zstd decompression streams cannot be seeked with SEEK_END'):

                reader.seek(0, os.SEEK_END)

            reader.close()

            with self.assertRaisesRegexp(ValueError, 'stream is closed'):

                reader.seek(4, os.SEEK_SET)

        with self.assertRaisesRegexp(ValueError, 'stream is closed'):

            reader.seek(0)

    def test_seek(self):

        source = b'foobar' * 60

        cctx = zstd.ZstdCompressor()

        frame = cctx.compress(source)

        dctx = zstd.ZstdDecompressor()

        with dctx.stream_reader(frame) as reader:

            reader.seek(3)

            self.assertEqual(reader.read(3), b'bar')

            reader.seek(4, os.SEEK_CUR)

            self.assertEqual(reader.read(2), b'ar')