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,

    TestCase,

)

if sys.version_info[0] >= 3:

    next = lambda it: it.__next__()

else:

    next = lambda it: it.next()

@make_cffi

class TestFrameHeaderSize(TestCase):

    def test_empty(self):

        with self.assertRaisesRegex(

            zstd.ZstdError,

            "could not determine frame header size: Src size " "is incorrect",

        ):

            zstd.frame_header_size(b"")

    def test_too_small(self):

        with self.assertRaisesRegex(

            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(TestCase):

    def test_empty(self):

        with self.assertRaisesRegex(

            zstd.ZstdError, "error when determining content size"

        ):

            zstd.frame_content_size(b"")

    def test_too_small(self):

        with self.assertRaisesRegex(

            zstd.ZstdError, "error when determining content size"

        ):

            zstd.frame_content_size(b"foob")

    def test_bad_frame(self):

        with self.assertRaisesRegex(

            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(TestCase):

    def test_memory_size(self):

        dctx = zstd.ZstdDecompressor()

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

@make_cffi

class TestDecompressor_decompress(TestCase):

    def test_empty_input(self):

        dctx = zstd.ZstdDecompressor()

        with self.assertRaisesRegex(

            zstd.ZstdError, "error determining content size from frame header"

        ):

            dctx.decompress(b"")

    def test_invalid_input(self):

        dctx = zstd.ZstdDecompressor()

        with self.assertRaisesRegex(

            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.assertRaisesRegex(

            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.assertRaisesRegex(

            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.assertRaisesRegex(

            zstd.ZstdError,

            "decompression error: Frame requires too much memory",

        ):

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

@make_cffi

class TestDecompressor_copy_stream(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(TestCase):

    def test_context_manager(self):

        dctx = zstd.ZstdDecompressor()

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

            with self.assertRaisesRegex(

                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.assertRaisesRegex(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.assertRaisesRegex(

                ValueError, "cannot read negative amounts less than -1"

            ):

                reader.read(-2)