#!/usr/bin/env python3
# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# Licensed 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.

# pyre-unsafe

import array
import struct
import sys
import unittest

from folly.iobuf import IOBuf, WritableIOBuf

from .iobuf_helper import (
    get_empty_chain,
    get_empty_writable_chain,
    make_chain,
    to_uppercase_string,
    to_uppercase_string_heap,
)


class IOBufTests(unittest.TestCase):
    def test_empty_chain(self) -> None:
        ebuf = get_empty_chain()
        self.assertFalse(ebuf)
        self.assertTrue(ebuf.is_chained)
        self.assertEqual(len(ebuf), 0)
        self.assertEqual(ebuf.chain_size(), 0)
        self.assertEqual(ebuf.chain_count(), 8)
        self.assertEqual(b"".join(ebuf), b"")
        self.assertEqual(b"", bytes(ebuf))

    def test_chain(self) -> None:
        control = [b"facebook", b"thrift", b"python3", b"cython"]
        chain = make_chain([IOBuf(x) for x in control])
        self.assertTrue(chain.is_chained)
        self.assertTrue(chain)
        self.assertEqual(bytes(chain), control[0])
        self.assertEqual(len(chain), len(control[0]))
        self.assertEqual(chain.chain_size(), sum(len(x) for x in control))
        self.assertEqual(chain.chain_count(), len(control))
        # pyre-fixme[6]: For 1st argument expected `Buffer` but got `Optional[IOBuf]`.
        self.assertEqual(memoryview(chain.next), control[1])
        self.assertEqual(b"".join(chain), b"".join(control))

    def test_cyclic_chain(self) -> None:
        control = [b"aaa", b"aaaa"]
        chain = make_chain([IOBuf(x) for x in control])
        self.assertTrue(chain.is_chained)
        self.assertTrue(chain)
        self.assertEqual(bytes(chain), control[0])
        self.assertEqual(len(chain), len(control[0]))
        self.assertEqual(chain.chain_size(), sum(len(x) for x in control))
        self.assertEqual(chain.chain_count(), len(control))
        # pyre-fixme[6]: For 1st argument expected `Buffer` but got `Optional[IOBuf]`.
        self.assertEqual(memoryview(chain.next), control[1])
        self.assertEqual(b"".join(chain), b"".join(control))

    def test_hash(self) -> None:
        x = b"omg"
        y = b"wtf"
        xb = IOBuf(x)
        yb = IOBuf(y)
        hash(xb)
        self.assertNotEqual(hash(xb), hash(yb))
        self.assertEqual(hash(xb), hash(IOBuf(x)))

    def test_empty(self) -> None:
        x = b""
        xb = IOBuf(x)
        self.assertEqual(memoryview(xb), x)
        self.assertEqual(bytes(xb), x)
        self.assertFalse(xb)
        self.assertEqual(len(xb), len(x))

    def test_iter(self) -> None:
        x = b"testtest"
        xb = IOBuf(x)
        self.assertEqual(b"".join(iter(xb)), x)

    def test_bytes(self) -> None:
        x = b"omgwtfbbq"
        xb = IOBuf(x)
        self.assertEqual(bytes(xb), x)

    def test_cmp(self) -> None:
        x = IOBuf(b"abc")
        y = IOBuf(b"def")
        z = IOBuf(b"abc")
        self.assertEqual(x, z)
        self.assertNotEqual(x, y)
        self.assertLess(x, y)
        self.assertLessEqual(x, y)
        self.assertLessEqual(x, z)
        self.assertGreater(y, x)
        self.assertGreaterEqual(y, x)

    def test_typed(self) -> None:
        # pyre-fixme[6]: Expected `Union[IOBuf, bytearray, bytes, memoryview]` for
        #  1st param but got `array[int]`.
        x = IOBuf(array.array("l", [1, 2, 3, 4, 5]))
        self.assertEqual(x.chain_size(), 5 * struct.calcsize("l"))

    def test_unshaped(self) -> None:
        x = IOBuf(memoryview(b"a").cast("B", shape=[]))
        self.assertEqual(x.chain_size(), 1)

    def test_multidimensional(self) -> None:
        x = IOBuf(memoryview(b"abcdef").cast("B", shape=[3, 2]))
        self.assertEqual(x.chain_size(), 6)

    def test_conversion_from_python_to_cpp(self) -> None:
        iobuf = make_chain(
            [
                IOBuf(memoryview(b"abc")),
                IOBuf(memoryview(b"def")),
                IOBuf(memoryview(b"ghi")),
            ]
        )
        self.assertEqual(to_uppercase_string(iobuf), "ABCDEFGHI")
        self.assertEqual(to_uppercase_string_heap(iobuf), "ABCDEFGHI")

    def test_conversion_from_python_to_cpp_with_wrong_type(self) -> None:
        not_an_iobuf = [1, 2, 3]
        with self.assertRaises(TypeError):
            to_uppercase_string(not_an_iobuf)
        with self.assertRaises(TypeError):
            to_uppercase_string_heap(not_an_iobuf)

    def test_buffer_read_fail(self) -> None:
        finish = bytearray(b"1234567890123456")
        xb = IOBuf(finish)
        buf = memoryview(xb)

        try:
            self.assertEqual(bytes(buf[17]), None)
            self.fail("Expected exception for reading out of bounds")
        except IndexError as e:
            self.assertEqual(
                str(e),
                "index out of bounds on dimension 1",
            )


class WritableIOBufTests(unittest.TestCase):
    def test_bytes_writable(self) -> None:
        x = bytearray(b"omgwtfbbq")
        xb = WritableIOBuf(x)
        self.assertEqual(bytes(xb), x)
        self.assertEqual(xb.writable(), True)

    def test_buffer_overwrite(self) -> None:
        start = b"omgwtfbbq"
        finish = b"123456789"
        x = bytearray(start)
        xb = WritableIOBuf(x)
        buf = memoryview(xb)

        self.assertEqual(bytes(buf), start)
        buf[:] = finish
        self.assertEqual(bytes(buf), finish)

        self.assertNotEqual(bytes(buf), start)

    def test_buffer_write_empty(self) -> None:
        start = bytearray(9)
        finish = b"123456789"
        xb = WritableIOBuf(bytearray(9))
        buf = memoryview(xb)

        self.assertEqual(bytes(buf), start)
        buf[:] = finish
        self.assertEqual(bytes(buf), finish)
        self.assertEqual(memoryview(xb), finish)

        self.assertNotEqual(bytes(buf), start)

    def test_buffer_update_in_place(self) -> None:
        x = bytearray(b"123")
        xb = WritableIOBuf(x)
        memoryview(xb)[:] = b"456"
        self.assertEqual(x, b"456")  # xb wrapped x, so by mutating xb we mutated x

    def test_buffer_write_out_of_bounds(self) -> None:
        start = bytearray(9)
        finish = b"1234567890"
        xb = WritableIOBuf(bytearray(9))
        buf = memoryview(xb)

        self.assertEqual(bytes(buf), start)
        try:
            buf[:] = finish
            self.fail("Expected exception for writing out of bounds")
        except ValueError as e:
            self.assertEqual(
                str(e),
                "memoryview assignment: lvalue and rvalue have different structures",
            )
        self.assertEqual(bytes(buf), start)

    def test_buffer_read_out_of_bounds(self) -> None:
        start = bytearray(b"1234567890")
        xb = WritableIOBuf(start)
        buf = memoryview(xb)

        self.assertEqual(bytes(buf), start)
        try:
            self.assertEqual(bytes(buf[10]), None)
            self.fail("Expected exception for reading out of bounds")
        except IndexError as e:
            self.assertEqual(
                str(e),
                "index out of bounds on dimension 1",
            )

    def test_buffer_write_empty_pieces(self) -> None:
        start = bytearray(9)
        builder = bytearray(9)
        piece_1 = b"123"
        piece_2 = b"456"
        piece_3 = b"789"
        xb = WritableIOBuf(start)

        buf = memoryview(xb)
        buf[0:3] = piece_1
        for i in range(3):
            builder[i] = piece_1[i]
        self.assertEqual(bytes(buf), builder)

        buf = memoryview(xb)
        buf[3:6] = piece_2
        for i in range(3):
            builder[i + 3] = piece_2[i]
        self.assertEqual(bytes(buf), builder)

        buf = memoryview(xb)
        buf[6:9] = piece_3
        for i in range(3):
            builder[i + 6] = piece_3[i]
        self.assertEqual(bytes(buf), builder)

    def test_empty_writable_chain(self) -> None:
        ebuf = get_empty_writable_chain()
        self.assertFalse(ebuf)
        self.assertTrue(ebuf.is_chained)
        self.assertEqual(len(ebuf), 0)
        self.assertEqual(ebuf.chain_size(), 0)
        self.assertEqual(ebuf.chain_count(), 8)
        self.assertEqual(b"".join(ebuf), b"")
        self.assertEqual(b"", bytes(ebuf))

    def test_appendable_writable_chain(self) -> None:
        x = bytearray(b"omgwtfbbq")
        xb = WritableIOBuf(x)
        self.assertFalse(xb.is_chained)

        y = bytearray(b"wtfbbqomg")
        yb = WritableIOBuf(y)
        xb.append_to_chain(yb)

        self.assertTrue(xb.is_chained)
        self.assertEqual(len(xb), 9)
        self.assertEqual(xb.chain_size(), 18)
        self.assertEqual(xb.chain_count(), 2)
        self.assertEqual(b"".join(xb), b"".join([x, y]))

    def test_appendable_writable_chain_overwrite(self) -> None:
        start = bytearray(9)
        x = bytearray(b"omgwtfbbq")
        for i in range(9):
            start[i] = x[i]

        xb = WritableIOBuf(x)
        self.assertFalse(xb.is_chained)

        y = bytearray(b"wtfbbqomg")
        yb = WritableIOBuf(y)
        xb.append_to_chain(yb)

        finish = b"123456789"
        buf = memoryview(xb)
        buf[:] = finish

        self.assertTrue(xb.is_chained)
        self.assertEqual(len(xb), 9)
        self.assertEqual(xb.chain_size(), 18)
        self.assertEqual(xb.chain_count(), 2)
        self.assertEqual(b"".join(xb), b"".join([x, y]))
        self.assertEqual(bytes(buf), x)

        self.assertNotEqual(b"".join(xb), b"".join([start, y]))

    def test_appendable_writable_chain_coalesce(self) -> None:
        x = bytearray(b"omgwtfbbq")
        xb = WritableIOBuf(x)
        self.assertFalse(xb.is_chained)

        y = bytearray(b"wtfbbqomg")
        yb = WritableIOBuf(y)
        xb.append_to_chain(yb)

        xb.coalesce()
        self.assertFalse(xb.is_chained)
        self.assertEqual(len(xb), 18)
        self.assertEqual(xb.chain_size(), 18)
        self.assertEqual(xb.chain_count(), 1)

        buf = memoryview(xb)
        self.assertEqual(bytes(buf), b"".join([x, y]))

    def test_appendable_writable_chain_coalesce_exception(self) -> None:
        x = bytearray(b"omgwtfbbq")
        xb = WritableIOBuf(x)
        self.assertFalse(xb.is_chained)

        y = bytearray(b"wtfbbqomg")
        yb = WritableIOBuf(y)
        xb.append_to_chain(yb)
        test = xb

        self.assertEqual(test, xb)
        self.assertEqual(sys.getrefcount(xb), 3)
        try:
            xb.coalesce()
            self.fail("Expected exception for too many references")
        except RuntimeError as e:
            self.assertEqual(
                str(e), "Cannot coalesce IOBuf with more than one reference"
            )

    def test_buffer_creation_with_size_write(self) -> None:
        finish = bytearray(b"123456789012345")
        xb = WritableIOBuf.create_unitialized(17)
        buf = memoryview(xb)

        self.assertEqual(len(xb), 0)
        self.assertEqual(len(buf), 0)
        self.assertEqual(bytes(buf), b"")
        self.assertEqual(xb.length(), 0)
        xb.append(15)
        buf = memoryview(xb)
        self.assertNotEqual(bytes(buf), b"")
        buf[0:15] = finish
        self.assertEqual(len(xb), 15)
        self.assertEqual(len(memoryview(xb)), 15)
        self.assertEqual(memoryview(xb), finish)
        self.assertEqual(buf, finish)
        self.assertEqual(xb.length(), 15)
        self.assertEqual(bytes(xb), b"123456789012345")

    def test_buffer_creation_with_size_write_fail(self) -> None:
        finish = bytearray(b"12345678901234567")
        xb = WritableIOBuf.create_unitialized(16)
        xb.append(16)
        buf = memoryview(xb)

        self.assertEqual(xb.capacity(), 16)
        try:
            buf[:] = finish
            self.fail("Expected exception for writing out of bounds")
        except ValueError as e:
            self.assertEqual(
                str(e),
                "memoryview assignment: lvalue and rvalue have different structures",
            )

    def test_buffer_creation_with_size_read_fail(self) -> None:
        finish = bytearray(b"1234567890123456")
        xb = WritableIOBuf.create_unitialized(16)
        xb.append(16)
        buf = memoryview(xb)

        buf[:] = finish
        try:
            self.assertEqual(bytes(buf[17]), None)
            self.fail("Expected exception for reading out of bounds")
        except IndexError as e:
            self.assertEqual(
                str(e),
                "index out of bounds on dimension 1",
            )

    def test_buffer_creation_with_size_append(self) -> None:
        xb = WritableIOBuf.create_unitialized(16)
        self.assertEqual(xb.length(), 0)
        self.assertEqual(xb.capacity(), 16)
        xb.append(16)
        self.assertEqual(xb.length(), 16)

    def test_buffer_creation_with_size_append_fail(self) -> None:
        xb = WritableIOBuf.create_unitialized(16)
        self.assertEqual(xb.length(), 0)
        self.assertEqual(xb.capacity(), 16)
        try:
            xb.append(17)
            self.fail("Expected exception for appending more than capacity")
        except ValueError as e:
            self.assertEqual(str(e), "Cannot append more than capacity")

    def test_buffer_creation_with_size_append_fail_negative(self) -> None:
        xb = WritableIOBuf.create_unitialized(16)
        self.assertEqual(xb.length(), 0)
        self.assertEqual(xb.capacity(), 16)
        try:
            xb.append(-1)
            self.fail("Expected exception for negative amount")
        except ValueError as e:
            self.assertEqual(str(e), "Cannot append, amount must be positive")

    def test_buffer_creation_with_size_multiple_append(self) -> None:
        xb = WritableIOBuf.create_unitialized(16)
        self.assertEqual(xb.length(), 0)
        self.assertEqual(xb.capacity(), 16)
        xb.append(15)
        xb.append(1)
        self.assertEqual(xb.length(), 16)

    def test_buffer_creation_with_size_multiple_append_fail(self) -> None:
        xb = WritableIOBuf.create_unitialized(16)
        self.assertEqual(xb.length(), 0)
        self.assertEqual(xb.capacity(), 16)
        xb.append(16)
        try:
            xb.append(1)
            self.fail("Expected exception for appending more than capacity")
        except ValueError as e:
            self.assertEqual(str(e), "Cannot append more than capacity")

    def test_buffer_creation_with_size_trim_start(self) -> None:
        xb = WritableIOBuf.create_unitialized(16)
        self.assertEqual(xb.length(), 0)
        self.assertEqual(xb.capacity(), 16)
        xb.append(5)
        self.assertEqual(xb.length(), 5)
        xb.trim_start(5)
        self.assertEqual(xb.length(), 0)

    def test_buffer_creation_with_size_write_trim_start(self) -> None:
        finish = bytearray(b"1234567890123456")
        xb = WritableIOBuf.create_unitialized(16)
        xb.append(16)
        buf = memoryview(xb)

        buf[0:16] = finish
        self.assertEqual(memoryview(xb), finish)
        self.assertEqual(buf, finish)
        self.assertEqual(xb.length(), 16)
        self.assertEqual(bytes(xb), b"1234567890123456")
        xb.trim_start(6)
        self.assertEqual(bytes(xb), b"7890123456")
        self.assertEqual(xb.length(), 10)
        xb.trim_start(6)
        self.assertEqual(bytes(xb), b"3456")
        self.assertEqual(xb.length(), 4)

    def test_buffer_creation_with_size_trim_start_fail(self) -> None:
        xb = WritableIOBuf.create_unitialized(16)
        self.assertEqual(xb.length(), 0)
        self.assertEqual(xb.capacity(), 16)
        xb.append(5)
        self.assertEqual(xb.length(), 5)
        try:
            xb.trim_start(6)
            self.fail("Expected exception for trimming more than length")
        except ValueError as e:
            self.assertEqual(str(e), "Cannot trim more than length")

    def test_buffer_creation_with_size_trim_start_multiple(self) -> None:
        xb = WritableIOBuf.create_unitialized(16)
        self.assertEqual(xb.length(), 0)
        self.assertEqual(xb.capacity(), 16)
        xb.append(5)
        self.assertEqual(xb.length(), 5)
        xb.trim_start(4)
        xb.trim_start(1)
        self.assertEqual(xb.length(), 0)

    def test_buffer_creation_with_size_trim_start_fail_multiple(self) -> None:
        xb = WritableIOBuf.create_unitialized(16)
        self.assertEqual(xb.length(), 0)
        self.assertEqual(xb.capacity(), 16)
        xb.append(5)
        self.assertEqual(xb.length(), 5)
        xb.trim_start(5)
        try:
            xb.trim_start(1)
            self.fail("Expected exception for trimming more than length")
        except ValueError as e:
            self.assertEqual(str(e), "Cannot trim more than length")

    def test_buffer_creation_with_size_trim_start_fail_negative(self) -> None:
        xb = WritableIOBuf.create_unitialized(16)
        self.assertEqual(xb.length(), 0)
        self.assertEqual(xb.capacity(), 16)

        try:
            xb.trim_start(-1)
            self.fail("Expected exception for negative amount")
        except ValueError as e:
            self.assertEqual(str(e), "Cannot trim start, amount must be positive")

    def test_buffer_creation_with_size_trim_end(self) -> None:
        xb = WritableIOBuf.create_unitialized(16)
        self.assertEqual(xb.length(), 0)
        self.assertEqual(xb.capacity(), 16)
        xb.append(5)
        self.assertEqual(xb.length(), 5)
        xb.trim_end(5)
        self.assertEqual(xb.length(), 0)

    def test_buffer_creation_with_size_write_trim_end(self) -> None:
        finish = bytearray(b"1234567890123456")
        xb = WritableIOBuf.create_unitialized(16)
        xb.append(16)
        buf = memoryview(xb)

        buf[0:16] = finish
        self.assertEqual(memoryview(xb), finish)
        self.assertEqual(buf, finish)
        self.assertEqual(xb.length(), 16)
        self.assertEqual(bytes(xb), b"1234567890123456")
        xb.trim_end(6)
        self.assertEqual(bytes(xb), b"1234567890")
        self.assertEqual(xb.length(), 10)
        xb.trim_end(6)
        self.assertEqual(bytes(xb), b"1234")
        self.assertEqual(xb.length(), 4)

    def test_buffer_creation_with_size_trim_end_fail(self) -> None:
        xb = WritableIOBuf.create_unitialized(16)
        self.assertEqual(xb.length(), 0)
        self.assertEqual(xb.capacity(), 16)
        xb.append(5)
        self.assertEqual(xb.length(), 5)
        try:
            xb.trim_end(6)
            self.fail("Expected exception for trimming more than length")
        except ValueError as e:
            self.assertEqual(str(e), "Cannot trim more than length")

    def test_buffer_creation_with_size_trim_end_multiple(self) -> None:
        xb = WritableIOBuf.create_unitialized(16)
        self.assertEqual(xb.length(), 0)
        self.assertEqual(xb.capacity(), 16)
        xb.append(5)
        self.assertEqual(xb.length(), 5)
        xb.trim_end(4)
        xb.trim_end(1)
        self.assertEqual(xb.length(), 0)

    def test_buffer_creation_with_size_trim_end_fail_multiple(self) -> None:
        xb = WritableIOBuf.create_unitialized(16)
        self.assertEqual(xb.length(), 0)
        self.assertEqual(xb.capacity(), 16)
        xb.append(5)
        self.assertEqual(xb.length(), 5)
        xb.trim_end(5)
        try:
            xb.trim_end(1)
            self.fail("Expected exception for trimming more than length")
        except ValueError as e:
            self.assertEqual(str(e), "Cannot trim more than length")

    def test_buffer_creation_with_size_trim_end_fail_negative(self) -> None:
        xb = WritableIOBuf.create_unitialized(16)
        self.assertEqual(xb.length(), 0)
        self.assertEqual(xb.capacity(), 16)

        try:
            xb.trim_end(-1)
            self.fail("Expected exception for negative amount")
        except ValueError as e:
            self.assertEqual(str(e), "Cannot trim end, amount must be positive")

    def test_buffer_creation_with_size_prepend(self) -> None:
        xb = WritableIOBuf.create_unitialized(16)
        self.assertEqual(xb.length(), 0)
        self.assertEqual(xb.capacity(), 16)
        xb.append(5)
        self.assertEqual(xb.length(), 5)
        xb.trim_start(5)
        self.assertEqual(xb.length(), 0)
        xb.prepend(5)
        self.assertEqual(xb.length(), 5)

    def test_buffer_creation_with_size_write_prepend(self) -> None:
        finish = bytearray(b"1234567890123456")
        xb = WritableIOBuf.create_unitialized(16)
        xb.append(16)
        buf = memoryview(xb)

        buf[0:16] = finish
        self.assertEqual(memoryview(xb), finish)
        self.assertEqual(buf, finish)
        self.assertEqual(xb.length(), 16)
        self.assertEqual(bytes(xb), b"1234567890123456")
        xb.trim_start(6)
        self.assertEqual(bytes(xb), b"7890123456")
        self.assertEqual(xb.length(), 10)
        xb.prepend(6)
        self.assertEqual(bytes(xb), b"1234567890123456")
        self.assertEqual(xb.length(), 16)

    def test_buffer_creation_with_size_prepend_fail(self) -> None:
        xb = WritableIOBuf.create_unitialized(16)
        self.assertEqual(xb.length(), 0)
        self.assertEqual(xb.capacity(), 16)
        xb.append(5)
        self.assertEqual(xb.length(), 5)
        xb.trim_start(5)
        self.assertEqual(xb.length(), 0)
        try:
            xb.prepend(6)
            self.fail("Expected exception for prepending more than headroom")
        except ValueError as e:
            self.assertEqual(str(e), "Cannot prepend more than headroom")

    def test_buffer_creation_with_size_prepend_fail_multiple(self) -> None:
        xb = WritableIOBuf.create_unitialized(16)
        self.assertEqual(xb.length(), 0)
        self.assertEqual(xb.capacity(), 16)
        xb.append(5)
        self.assertEqual(xb.length(), 5)
        xb.trim_start(5)
        self.assertEqual(xb.length(), 0)
        xb.prepend(5)
        try:
            xb.prepend(1)
            self.fail("Expected exception for prepending more than headroom")
        except ValueError as e:
            self.assertEqual(str(e), "Cannot prepend more than headroom")

    def test_buffer_creation_with_size_prepend_fail_negative(self) -> None:
        xb = WritableIOBuf.create_unitialized(16)
        self.assertEqual(xb.length(), 0)
        self.assertEqual(xb.capacity(), 16)

        try:
            xb.prepend(-1)
            self.fail("Expected exception for negative amount")
        except ValueError as e:
            self.assertEqual(str(e), "Cannot prepend, amount must be positive")
