import io
import unittest

import zstandard as zstd

from .common import (
    NonClosingBytesIO,
    CustomBytesIO,
)


class TestCompressor_stream_reader(unittest.TestCase):
    def test_context_manager(self):
        cctx = zstd.ZstdCompressor()

        with cctx.stream_reader(b"foo") as reader:
            with self.assertRaisesRegex(
                ValueError, "cannot __enter__ multiple times"
            ):
                with reader as reader2:
                    pass

    def test_no_context_manager(self):
        cctx = zstd.ZstdCompressor()

        reader = cctx.stream_reader(b"foo")
        reader.read(4)
        self.assertFalse(reader.closed)

        reader.close()
        self.assertTrue(reader.closed)
        with self.assertRaisesRegex(ValueError, "stream is closed"):
            reader.read(1)

    def test_not_implemented(self):
        cctx = zstd.ZstdCompressor()

        with cctx.stream_reader(b"foo" * 60) 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(OSError):
                reader.writelines([])

            with self.assertRaises(OSError):
                reader.write(b"foo")

    def test_constant_methods(self):
        cctx = zstd.ZstdCompressor()

        with cctx.stream_reader(b"boo") as reader:
            self.assertTrue(reader.readable())
            self.assertFalse(reader.writable())
            self.assertFalse(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):
        cctx = zstd.ZstdCompressor()

        with cctx.stream_reader(b"foo" * 60) as reader:
            reader.close()
            self.assertTrue(reader.closed)
            with self.assertRaisesRegex(ValueError, "stream is closed"):
                reader.read(10)

    def test_read_sizes(self):
        cctx = zstd.ZstdCompressor()
        foo = cctx.compress(b"foo")

        with cctx.stream_reader(b"foo") as reader:
            with self.assertRaisesRegex(
                ValueError, "cannot read negative amounts less than -1"
            ):
                reader.read(-2)

            self.assertEqual(reader.read(0), b"")
            self.assertEqual(reader.read(), 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)

        with cctx.stream_reader(source) as reader:
            self.assertEqual(reader.tell(), 0)

            # We should get entire frame in one read.
            result = reader.read(8192)
            self.assertEqual(result, frame)
            self.assertEqual(reader.tell(), len(result))
            self.assertEqual(reader.read(), b"")
            self.assertEqual(reader.tell(), len(result))

    def test_read_buffer_small_chunks(self):
        cctx = zstd.ZstdCompressor()

        source = b"foo" * 60
        chunks = []

        with cctx.stream_reader(source) as reader:
            self.assertEqual(reader.tell(), 0)

            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), cctx.compress(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)

        with cctx.stream_reader(io.BytesIO(source), size=len(source)) as reader:
            self.assertEqual(reader.tell(), 0)

            chunk = reader.read(8192)
            self.assertEqual(chunk, frame)
            self.assertEqual(reader.tell(), len(chunk))
            self.assertEqual(reader.read(), b"")
            self.assertEqual(reader.tell(), len(chunk))

    def test_read_stream_small_chunks(self):
        cctx = zstd.ZstdCompressor()

        source = b"foo" * 60
        chunks = []

        with cctx.stream_reader(io.BytesIO(source), size=len(source)) as reader:
            self.assertEqual(reader.tell(), 0)

            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), cctx.compress(source))

    def test_read_after_exit(self):
        cctx = zstd.ZstdCompressor()

        with cctx.stream_reader(b"foo" * 60) as reader:
            while reader.read(8192):
                pass

        with self.assertRaisesRegex(ValueError, "stream is closed"):
            reader.read(10)

    def test_bad_size(self):
        cctx = zstd.ZstdCompressor()

        source = io.BytesIO(b"foobar")

        with cctx.stream_reader(source, size=2) as reader:
            with self.assertRaisesRegex(
                zstd.ZstdError, "Src size is incorrect"
            ):
                reader.read(10)

        # Try another compression operation.
        with cctx.stream_reader(source, size=42):
            pass

    def test_readall(self):
        cctx = zstd.ZstdCompressor()
        frame = cctx.compress(b"foo" * 1024)

        reader = cctx.stream_reader(b"foo" * 1024)
        self.assertEqual(reader.readall(), frame)

    def test_readinto(self):
        cctx = zstd.ZstdCompressor()
        foo = cctx.compress(b"foo")

        reader = cctx.stream_reader(b"foo")
        with self.assertRaises(Exception):
            reader.readinto(b"foobar")

        # readinto() with sufficiently large destination.
        b = bytearray(1024)
        reader = cctx.stream_reader(b"foo")
        self.assertEqual(reader.readinto(b), len(foo))
        self.assertEqual(b[0 : len(foo)], foo)
        self.assertEqual(reader.readinto(b), 0)
        self.assertEqual(b[0 : len(foo)], foo)

        # readinto() with small reads.
        b = bytearray(1024)
        reader = cctx.stream_reader(b"foo", read_size=1)
        self.assertEqual(reader.readinto(b), len(foo))
        self.assertEqual(b[0 : len(foo)], foo)

        # Too small destination buffer.
        b = bytearray(2)
        reader = cctx.stream_reader(b"foo")
        self.assertEqual(reader.readinto(b), 2)
        self.assertEqual(b[:], foo[0:2])
        self.assertEqual(reader.readinto(b), 2)
        self.assertEqual(b[:], foo[2:4])
        self.assertEqual(reader.readinto(b), 2)
        self.assertEqual(b[:], foo[4:6])

    def test_readinto1(self):
        cctx = zstd.ZstdCompressor()
        foo = b"".join(cctx.read_to_iter(io.BytesIO(b"foo")))

        reader = cctx.stream_reader(b"foo")
        with self.assertRaises(Exception):
            reader.readinto1(b"foobar")

        b = bytearray(1024)
        source = CustomBytesIO(b"foo")
        reader = cctx.stream_reader(source)
        self.assertEqual(reader.readinto1(b), len(foo))
        self.assertEqual(b[0 : len(foo)], foo)
        self.assertEqual(source._read_count, 2)

        # readinto1() with small reads.
        b = bytearray(1024)
        source = CustomBytesIO(b"foo")
        reader = cctx.stream_reader(source, read_size=1)
        self.assertEqual(reader.readinto1(b), len(foo))
        self.assertEqual(b[0 : len(foo)], foo)
        self.assertEqual(source._read_count, 4)

    def test_read1(self):
        cctx = zstd.ZstdCompressor()
        foo = b"".join(cctx.read_to_iter(io.BytesIO(b"foo")))

        b = CustomBytesIO(b"foo")
        reader = cctx.stream_reader(b)

        self.assertEqual(reader.read1(), foo)
        self.assertEqual(b._read_count, 2)

        b = CustomBytesIO(b"foo")
        reader = cctx.stream_reader(b)

        self.assertEqual(reader.read1(0), b"")
        self.assertEqual(reader.read1(2), foo[0:2])
        self.assertEqual(b._read_count, 2)
        self.assertEqual(reader.read1(2), foo[2:4])
        self.assertEqual(reader.read1(1024), foo[4:])

    def test_close(self):
        buffer = NonClosingBytesIO(b"foo" * 1024)
        cctx = zstd.ZstdCompressor()
        reader = cctx.stream_reader(buffer)

        reader.read(3)
        self.assertFalse(reader.closed)
        self.assertFalse(buffer.closed)
        reader.close()
        self.assertTrue(reader.closed)
        self.assertTrue(buffer.closed)

        with self.assertRaisesRegex(ValueError, "stream is closed"):
            reader.read(3)

        with self.assertRaisesRegex(ValueError, "stream is closed"):
            with reader:
                pass

        # Context manager exit should close stream.
        buffer = io.BytesIO(b"foo" * 1024)
        reader = cctx.stream_reader(buffer)

        with reader:
            reader.read(3)

        self.assertTrue(reader.closed)
        self.assertTrue(buffer.closed)

        # Context manager exit should close stream if an exception raised.
        buffer = io.BytesIO(b"foo" * 1024)
        reader = cctx.stream_reader(buffer)

        with self.assertRaisesRegex(Exception, "ignore"):
            with reader:
                reader.read(3)
                raise Exception("ignore")

        self.assertTrue(reader.closed)
        self.assertTrue(buffer.closed)

        # Test with non-file source.
        with cctx.stream_reader(b"foo" * 1024) as reader:
            reader.read(3)
            self.assertFalse(reader.closed)

        self.assertTrue(reader.closed)

    def test_close_closefd_false(self):
        buffer = NonClosingBytesIO(b"foo" * 1024)
        cctx = zstd.ZstdCompressor()
        reader = cctx.stream_reader(buffer, closefd=False)

        reader.read(3)
        self.assertFalse(reader.closed)
        self.assertFalse(buffer.closed)
        reader.close()
        self.assertTrue(reader.closed)
        self.assertFalse(buffer.closed)

        with self.assertRaisesRegex(ValueError, "stream is closed"):
            reader.read(3)

        with self.assertRaisesRegex(ValueError, "stream is closed"):
            with reader:
                pass

        # Context manager exit should close stream.
        buffer = io.BytesIO(b"foo" * 1024)
        reader = cctx.stream_reader(buffer, closefd=False)

        with reader:
            reader.read(3)

        self.assertTrue(reader.closed)
        self.assertFalse(buffer.closed)

        # Context manager exit should close stream if an exception raised.
        buffer = io.BytesIO(b"foo" * 1024)
        reader = cctx.stream_reader(buffer, closefd=False)

        with self.assertRaisesRegex(Exception, "ignore"):
            with reader:
                reader.read(3)
                raise Exception("ignore")

        self.assertTrue(reader.closed)
        self.assertFalse(buffer.closed)

        # Test with non-file source variant.
        with cctx.stream_reader(b"foo" * 1024, closefd=False) as reader:
            reader.read(3)
            self.assertFalse(reader.closed)

        self.assertTrue(reader.closed)

    def test_write_exception(self):
        b = CustomBytesIO()
        b.write_exception = IOError("write")

        cctx = zstd.ZstdCompressor()

        writer = cctx.stream_writer(b)
        # Initial write won't issue write() to underlying stream.
        writer.write(b"foo")

        with self.assertRaisesRegex(IOError, "write"):
            writer.flush()