# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import unittest

from io import BytesIO

from cassandra import DriverException
from cassandra.segment import Segment, CrcException
from cassandra.connection import segment_codec_no_compression, segment_codec_lz4


def to_bits(b):
    return '{:08b}'.format(b)


class SegmentCodecTest(unittest.TestCase):

    small_msg = b'b' * 50
    max_msg = b'b' * Segment.MAX_PAYLOAD_LENGTH
    large_msg = b'b' * (Segment.MAX_PAYLOAD_LENGTH + 1)

    @staticmethod
    def _header_to_bits(data):
        # unpack a header to bits
        # data should be the little endian bytes sequence
        if len(data) > 6:  # compressed
            data = data[:5]
            bits = ''.join([to_bits(b) for b in reversed(data)])
            # return the compressed payload length, the uncompressed payload length,
            # the self-contained flag and the padding as bits
            return bits[23:40] + bits[6:23] + bits[5:6] + bits[:5]
        else:  # uncompressed
            data = data[:3]
            bits = ''.join([to_bits(b) for b in reversed(data)])
            # return the payload length, the self-contained flag and
            # the padding as bits
            return bits[7:24] + bits[6:7] + bits[:6]

    def test_encode_uncompressed_header(self):
        buffer = BytesIO()
        segment_codec_no_compression.encode_header(buffer, len(self.small_msg), -1, True)
        self.assertEqual(buffer.tell(), 6)
        self.assertEqual(
            self._header_to_bits(buffer.getvalue()),
            "00000000000110010" + "1" + "000000")

    @unittest.skipUnless(segment_codec_lz4, ' lz4 not installed')
    def test_encode_compressed_header(self):
        buffer = BytesIO()
        compressed_length = len(segment_codec_lz4.compress(self.small_msg))
        segment_codec_lz4.encode_header(buffer, compressed_length, len(self.small_msg), True)

        self.assertEqual(buffer.tell(), 8)
        self.assertEqual(
            self._header_to_bits(buffer.getvalue()),
            "{:017b}".format(compressed_length) + "00000000000110010" + "1" + "00000")

    def test_encode_uncompressed_header_with_max_payload(self):
        buffer = BytesIO()
        segment_codec_no_compression.encode_header(buffer, len(self.max_msg), -1, True)
        self.assertEqual(buffer.tell(), 6)
        self.assertEqual(
            self._header_to_bits(buffer.getvalue()),
            "11111111111111111" + "1" + "000000")

    def test_encode_header_fails_if_payload_too_big(self):
        buffer = BytesIO()
        for codec in [c for c in [segment_codec_no_compression, segment_codec_lz4] if c is not None]:
            with self.assertRaises(DriverException):
                codec.encode_header(buffer, len(self.large_msg), -1, False)

    def test_encode_uncompressed_header_not_self_contained_msg(self):
        buffer = BytesIO()
        # simulate the first chunk with the max size
        segment_codec_no_compression.encode_header(buffer, len(self.max_msg), -1, False)
        self.assertEqual(buffer.tell(), 6)
        self.assertEqual(
            self._header_to_bits(buffer.getvalue()),
            ("11111111111111111"
             "0"  # not self-contained
             "000000"))

    @unittest.skipUnless(segment_codec_lz4, ' lz4 not installed')
    def test_encode_compressed_header_with_max_payload(self):
        buffer = BytesIO()
        compressed_length = len(segment_codec_lz4.compress(self.max_msg))
        segment_codec_lz4.encode_header(buffer, compressed_length, len(self.max_msg), True)
        self.assertEqual(buffer.tell(), 8)
        self.assertEqual(
            self._header_to_bits(buffer.getvalue()),
            "{:017b}".format(compressed_length) + "11111111111111111" + "1" + "00000")

    @unittest.skipUnless(segment_codec_lz4, ' lz4 not installed')
    def test_encode_compressed_header_not_self_contained_msg(self):
        buffer = BytesIO()
        # simulate the first chunk with the max size
        compressed_length = len(segment_codec_lz4.compress(self.max_msg))
        segment_codec_lz4.encode_header(buffer, compressed_length, len(self.max_msg), False)
        self.assertEqual(buffer.tell(), 8)
        self.assertEqual(
            self._header_to_bits(buffer.getvalue()),
            ("{:017b}".format(compressed_length) +
             "11111111111111111"
             "0"  # not self-contained
             "00000"))

    def test_decode_uncompressed_header(self):
        buffer = BytesIO()
        segment_codec_no_compression.encode_header(buffer, len(self.small_msg), -1, True)
        buffer.seek(0)
        header = segment_codec_no_compression.decode_header(buffer)
        self.assertEqual(header.uncompressed_payload_length, -1)
        self.assertEqual(header.payload_length, len(self.small_msg))
        self.assertEqual(header.is_self_contained, True)

    @unittest.skipUnless(segment_codec_lz4, ' lz4 not installed')
    def test_decode_compressed_header(self):
        buffer = BytesIO()
        compressed_length = len(segment_codec_lz4.compress(self.small_msg))
        segment_codec_lz4.encode_header(buffer, compressed_length, len(self.small_msg), True)
        buffer.seek(0)
        header = segment_codec_lz4.decode_header(buffer)
        self.assertEqual(header.uncompressed_payload_length, len(self.small_msg))
        self.assertEqual(header.payload_length, compressed_length)
        self.assertEqual(header.is_self_contained, True)

    def test_decode_header_fails_if_corrupted(self):
        buffer = BytesIO()
        segment_codec_no_compression.encode_header(buffer, len(self.small_msg), -1, True)
        # corrupt one byte
        buffer.seek(buffer.tell()-1)
        buffer.write(b'0')
        buffer.seek(0)

        with self.assertRaises(CrcException):
            segment_codec_no_compression.decode_header(buffer)

    def test_decode_uncompressed_self_contained_segment(self):
        buffer = BytesIO()
        segment_codec_no_compression.encode(buffer, self.small_msg)

        buffer.seek(0)
        header = segment_codec_no_compression.decode_header(buffer)
        segment = segment_codec_no_compression.decode(buffer, header)

        self.assertEqual(header.is_self_contained, True)
        self.assertEqual(header.uncompressed_payload_length, -1)
        self.assertEqual(header.payload_length, len(self.small_msg))
        self.assertEqual(segment.payload, self.small_msg)

    @unittest.skipUnless(segment_codec_lz4, ' lz4 not installed')
    def test_decode_compressed_self_contained_segment(self):
        buffer = BytesIO()
        segment_codec_lz4.encode(buffer, self.small_msg)

        buffer.seek(0)
        header = segment_codec_lz4.decode_header(buffer)
        segment = segment_codec_lz4.decode(buffer, header)

        self.assertEqual(header.is_self_contained, True)
        self.assertEqual(header.uncompressed_payload_length, len(self.small_msg))
        self.assertGreater(header.uncompressed_payload_length, header.payload_length)
        self.assertEqual(segment.payload, self.small_msg)

    def test_decode_multi_segments(self):
        buffer = BytesIO()
        segment_codec_no_compression.encode(buffer, self.large_msg)

        buffer.seek(0)
        # We should have 2 segments to read
        headers = []
        segments = []
        headers.append(segment_codec_no_compression.decode_header(buffer))
        segments.append(segment_codec_no_compression.decode(buffer, headers[0]))
        headers.append(segment_codec_no_compression.decode_header(buffer))
        segments.append(segment_codec_no_compression.decode(buffer, headers[1]))

        self.assertTrue(all([h.is_self_contained is False for h in headers]))
        decoded_msg = segments[0].payload + segments[1].payload
        self.assertEqual(decoded_msg, self.large_msg)

    @unittest.skipUnless(segment_codec_lz4, ' lz4 not installed')
    def test_decode_fails_if_corrupted(self):
        buffer = BytesIO()
        segment_codec_lz4.encode(buffer, self.small_msg)
        buffer.seek(buffer.tell()-1)
        buffer.write(b'0')
        buffer.seek(0)
        header = segment_codec_lz4.decode_header(buffer)
        with self.assertRaises(CrcException):
            segment_codec_lz4.decode(buffer, header)

    @unittest.skipUnless(segment_codec_lz4, ' lz4 not installed')
    def test_decode_tiny_msg_not_compressed(self):
        buffer = BytesIO()
        segment_codec_lz4.encode(buffer, b'b')
        buffer.seek(0)
        header = segment_codec_lz4.decode_header(buffer)
        segment = segment_codec_lz4.decode(buffer, header)
        self.assertEqual(header.uncompressed_payload_length, 0)
        self.assertEqual(header.payload_length, 1)
        self.assertEqual(segment.payload, b'b')