import codecs
import dataclasses
import unittest
from unittest.mock import patch

from websockets.exceptions import PayloadTooBig, ProtocolError
from websockets.frames import *
from websockets.frames import CloseCode
from websockets.streams import StreamReader

from .utils import GeneratorTestCase


class FramesTestCase(GeneratorTestCase):
    def parse(self, data, mask, max_size=None, extensions=None):
        """
        Parse a frame from a bytestring.

        """
        reader = StreamReader()
        reader.feed_data(data)
        reader.feed_eof()
        parser = Frame.parse(
            reader.read_exact, mask=mask, max_size=max_size, extensions=extensions
        )
        return self.assertGeneratorReturns(parser)

    def assertFrameData(self, frame, data, mask, extensions=None):
        """
        Serializing frame yields data. Parsing data yields frame.

        """
        # Compare frames first, because test failures are easier to read,
        # especially when mask = True.
        parsed = self.parse(data, mask=mask, extensions=extensions)
        self.assertEqual(parsed, frame)

        # Make masking deterministic by reusing the same "random" mask.
        # This has an effect only when mask is True.
        mask_bytes = data[2:6] if mask else b""
        with patch("secrets.token_bytes", return_value=mask_bytes):
            serialized = frame.serialize(mask=mask, extensions=extensions)
        self.assertEqual(serialized, data)


class FrameTests(FramesTestCase):
    def test_text_unmasked(self):
        self.assertFrameData(
            Frame(OP_TEXT, b"Spam"),
            b"\x81\x04Spam",
            mask=False,
        )

    def test_text_masked(self):
        self.assertFrameData(
            Frame(OP_TEXT, b"Spam"),
            b"\x81\x84\x5b\xfb\xe1\xa8\x08\x8b\x80\xc5",
            mask=True,
        )

    def test_binary_unmasked(self):
        self.assertFrameData(
            Frame(OP_BINARY, b"Eggs"),
            b"\x82\x04Eggs",
            mask=False,
        )

    def test_binary_masked(self):
        self.assertFrameData(
            Frame(OP_BINARY, b"Eggs"),
            b"\x82\x84\x53\xcd\xe2\x89\x16\xaa\x85\xfa",
            mask=True,
        )

    def test_non_ascii_text_unmasked(self):
        self.assertFrameData(
            Frame(OP_TEXT, "café".encode()),
            b"\x81\x05caf\xc3\xa9",
            mask=False,
        )

    def test_non_ascii_text_masked(self):
        self.assertFrameData(
            Frame(OP_TEXT, "café".encode()),
            b"\x81\x85\x64\xbe\xee\x7e\x07\xdf\x88\xbd\xcd",
            mask=True,
        )

    def test_close(self):
        self.assertFrameData(
            Frame(OP_CLOSE, b""),
            b"\x88\x00",
            mask=False,
        )

    def test_ping(self):
        self.assertFrameData(
            Frame(OP_PING, b"ping"),
            b"\x89\x04ping",
            mask=False,
        )

    def test_pong(self):
        self.assertFrameData(
            Frame(OP_PONG, b"pong"),
            b"\x8a\x04pong",
            mask=False,
        )

    def test_long(self):
        self.assertFrameData(
            Frame(OP_BINARY, 126 * b"a"),
            b"\x82\x7e\x00\x7e" + 126 * b"a",
            mask=False,
        )

    def test_very_long(self):
        self.assertFrameData(
            Frame(OP_BINARY, 65536 * b"a"),
            b"\x82\x7f\x00\x00\x00\x00\x00\x01\x00\x00" + 65536 * b"a",
            mask=False,
        )

    def test_payload_too_big(self):
        with self.assertRaises(PayloadTooBig):
            self.parse(b"\x82\x7e\x04\x01" + 1025 * b"a", mask=False, max_size=1024)

    def test_bad_reserved_bits(self):
        for data in [b"\xc0\x00", b"\xa0\x00", b"\x90\x00"]:
            with self.subTest(data=data):
                with self.assertRaises(ProtocolError):
                    self.parse(data, mask=False)

    def test_good_opcode(self):
        for opcode in list(range(0x00, 0x03)) + list(range(0x08, 0x0B)):
            data = bytes([0x80 | opcode, 0])
            with self.subTest(data=data):
                self.parse(data, mask=False)  # does not raise an exception

    def test_bad_opcode(self):
        for opcode in list(range(0x03, 0x08)) + list(range(0x0B, 0x10)):
            data = bytes([0x80 | opcode, 0])
            with self.subTest(data=data):
                with self.assertRaises(ProtocolError):
                    self.parse(data, mask=False)

    def test_mask_flag(self):
        # Mask flag correctly set.
        self.parse(b"\x80\x80\x00\x00\x00\x00", mask=True)
        # Mask flag incorrectly unset.
        with self.assertRaises(ProtocolError):
            self.parse(b"\x80\x80\x00\x00\x00\x00", mask=False)
        # Mask flag correctly unset.
        self.parse(b"\x80\x00", mask=False)
        # Mask flag incorrectly set.
        with self.assertRaises(ProtocolError):
            self.parse(b"\x80\x00", mask=True)

    def test_control_frame_max_length(self):
        # At maximum allowed length.
        self.parse(b"\x88\x7e\x00\x7d" + 125 * b"a", mask=False)
        # Above maximum allowed length.
        with self.assertRaises(ProtocolError):
            self.parse(b"\x88\x7e\x00\x7e" + 126 * b"a", mask=False)

    def test_fragmented_control_frame(self):
        # Fin bit correctly set.
        self.parse(b"\x88\x00", mask=False)
        # Fin bit incorrectly unset.
        with self.assertRaises(ProtocolError):
            self.parse(b"\x08\x00", mask=False)

    def test_extensions(self):
        class Rot13:
            @staticmethod
            def encode(frame):
                assert frame.opcode == OP_TEXT
                text = frame.data.decode()
                data = codecs.encode(text, "rot13").encode()
                return dataclasses.replace(frame, data=data)

            # This extensions is symmetrical.
            @staticmethod
            def decode(frame, *, max_size=None):
                return Rot13.encode(frame)

        self.assertFrameData(
            Frame(OP_TEXT, b"hello"),
            b"\x81\x05uryyb",
            mask=False,
            extensions=[Rot13()],
        )


class StrTests(unittest.TestCase):
    def test_cont_text(self):
        self.assertEqual(
            str(Frame(OP_CONT, b" cr\xc3\xa8me", fin=False)),
            "CONT ' crème' [text, 7 bytes, continued]",
        )

    def test_cont_binary(self):
        self.assertEqual(
            str(Frame(OP_CONT, b"\xfc\xfd\xfe\xff", fin=False)),
            "CONT fc fd fe ff [binary, 4 bytes, continued]",
        )

    def test_cont_binary_from_memoryview(self):
        self.assertEqual(
            str(Frame(OP_CONT, memoryview(b"\xfc\xfd\xfe\xff"), fin=False)),
            "CONT fc fd fe ff [binary, 4 bytes, continued]",
        )

    def test_cont_final_text(self):
        self.assertEqual(
            str(Frame(OP_CONT, b" cr\xc3\xa8me")),
            "CONT ' crème' [text, 7 bytes]",
        )

    def test_cont_final_binary(self):
        self.assertEqual(
            str(Frame(OP_CONT, b"\xfc\xfd\xfe\xff")),
            "CONT fc fd fe ff [binary, 4 bytes]",
        )

    def test_cont_final_binary_from_memoryview(self):
        self.assertEqual(
            str(Frame(OP_CONT, memoryview(b"\xfc\xfd\xfe\xff"))),
            "CONT fc fd fe ff [binary, 4 bytes]",
        )

    def test_cont_text_truncated(self):
        self.assertEqual(
            str(Frame(OP_CONT, b"caf\xc3\xa9 " * 16, fin=False)),
            "CONT 'café café café café café café café café café ca..."
            "fé café café café café ' [text, 96 bytes, continued]",
        )

    def test_cont_binary_truncated(self):
        self.assertEqual(
            str(Frame(OP_CONT, bytes(range(256)), fin=False)),
            "CONT 00 01 02 03 04 05 06 07 08 09 0a 0b 0c 0d 0e 0f ..."
            " f8 f9 fa fb fc fd fe ff [binary, 256 bytes, continued]",
        )

    def test_cont_binary_truncated_from_memoryview(self):
        self.assertEqual(
            str(Frame(OP_CONT, memoryview(bytes(range(256))), fin=False)),
            "CONT 00 01 02 03 04 05 06 07 08 09 0a 0b 0c 0d 0e 0f ..."
            " f8 f9 fa fb fc fd fe ff [binary, 256 bytes, continued]",
        )

    def test_text(self):
        self.assertEqual(
            str(Frame(OP_TEXT, b"caf\xc3\xa9")),
            "TEXT 'café' [5 bytes]",
        )

    def test_text_non_final(self):
        self.assertEqual(
            str(Frame(OP_TEXT, b"caf\xc3\xa9", fin=False)),
            "TEXT 'café' [5 bytes, continued]",
        )

    def test_text_truncated(self):
        self.assertEqual(
            str(Frame(OP_TEXT, b"caf\xc3\xa9 " * 16)),
            "TEXT 'café café café café café café café café café ca..."
            "fé café café café café ' [96 bytes]",
        )

    def test_text_with_newline(self):
        self.assertEqual(
            str(Frame(OP_TEXT, b"Hello\nworld!")),
            "TEXT 'Hello\\nworld!' [12 bytes]",
        )

    def test_binary(self):
        self.assertEqual(
            str(Frame(OP_BINARY, b"\x00\x01\x02\x03")),
            "BINARY 00 01 02 03 [4 bytes]",
        )

    def test_binary_from_memoryview(self):
        self.assertEqual(
            str(Frame(OP_BINARY, memoryview(b"\x00\x01\x02\x03"))),
            "BINARY 00 01 02 03 [4 bytes]",
        )

    def test_binary_non_final(self):
        self.assertEqual(
            str(Frame(OP_BINARY, b"\x00\x01\x02\x03", fin=False)),
            "BINARY 00 01 02 03 [4 bytes, continued]",
        )

    def test_binary_non_final_from_memoryview(self):
        self.assertEqual(
            str(Frame(OP_BINARY, memoryview(b"\x00\x01\x02\x03"), fin=False)),
            "BINARY 00 01 02 03 [4 bytes, continued]",
        )

    def test_binary_truncated(self):
        self.assertEqual(
            str(Frame(OP_BINARY, bytes(range(256)))),
            "BINARY 00 01 02 03 04 05 06 07 08 09 0a 0b 0c 0d 0e 0f ..."
            " f8 f9 fa fb fc fd fe ff [256 bytes]",
        )

    def test_binary_truncated_from_memoryview(self):
        self.assertEqual(
            str(Frame(OP_BINARY, memoryview(bytes(range(256))))),
            "BINARY 00 01 02 03 04 05 06 07 08 09 0a 0b 0c 0d 0e 0f ..."
            " f8 f9 fa fb fc fd fe ff [256 bytes]",
        )

    def test_close(self):
        self.assertEqual(
            str(Frame(OP_CLOSE, b"\x03\xe8")),
            "CLOSE 1000 (OK) [2 bytes]",
        )

    def test_close_reason(self):
        self.assertEqual(
            str(Frame(OP_CLOSE, b"\x03\xe9Bye!")),
            "CLOSE 1001 (going away) Bye! [6 bytes]",
        )

    def test_ping(self):
        self.assertEqual(
            str(Frame(OP_PING, b"")),
            "PING '' [0 bytes]",
        )

    def test_ping_text(self):
        self.assertEqual(
            str(Frame(OP_PING, b"ping")),
            "PING 'ping' [text, 4 bytes]",
        )

    def test_ping_text_with_newline(self):
        self.assertEqual(
            str(Frame(OP_PING, b"ping\n")),
            "PING 'ping\\n' [text, 5 bytes]",
        )

    def test_ping_binary(self):
        self.assertEqual(
            str(Frame(OP_PING, b"\xff\x00\xff\x00")),
            "PING ff 00 ff 00 [binary, 4 bytes]",
        )

    def test_pong(self):
        self.assertEqual(
            str(Frame(OP_PONG, b"")),
            "PONG '' [0 bytes]",
        )

    def test_pong_text(self):
        self.assertEqual(
            str(Frame(OP_PONG, b"pong")),
            "PONG 'pong' [text, 4 bytes]",
        )

    def test_pong_text_with_newline(self):
        self.assertEqual(
            str(Frame(OP_PONG, b"pong\n")),
            "PONG 'pong\\n' [text, 5 bytes]",
        )

    def test_pong_binary(self):
        self.assertEqual(
            str(Frame(OP_PONG, b"\xff\x00\xff\x00")),
            "PONG ff 00 ff 00 [binary, 4 bytes]",
        )


class CloseTests(unittest.TestCase):
    def assertCloseData(self, close, data):
        """
        Serializing close yields data. Parsing data yields close.

        """
        serialized = close.serialize()
        self.assertEqual(serialized, data)
        parsed = Close.parse(data)
        self.assertEqual(parsed, close)

    def test_str(self):
        self.assertEqual(
            str(Close(CloseCode.NORMAL_CLOSURE, "")),
            "1000 (OK)",
        )
        self.assertEqual(
            str(Close(CloseCode.GOING_AWAY, "Bye!")),
            "1001 (going away) Bye!",
        )
        self.assertEqual(
            str(Close(3000, "")),
            "3000 (registered)",
        )
        self.assertEqual(
            str(Close(4000, "")),
            "4000 (private use)",
        )
        self.assertEqual(
            str(Close(5000, "")),
            "5000 (unknown)",
        )

    def test_parse_and_serialize(self):
        self.assertCloseData(
            Close(CloseCode.NORMAL_CLOSURE, "OK"),
            b"\x03\xe8OK",
        )
        self.assertCloseData(
            Close(CloseCode.GOING_AWAY, ""),
            b"\x03\xe9",
        )

    def test_parse_empty(self):
        self.assertEqual(
            Close.parse(b""),
            Close(CloseCode.NO_STATUS_RCVD, ""),
        )

    def test_parse_errors(self):
        with self.assertRaises(ProtocolError):
            Close.parse(b"\x03")
        with self.assertRaises(ProtocolError):
            Close.parse(b"\x03\xe7")
        with self.assertRaises(UnicodeDecodeError):
            Close.parse(b"\x03\xe8\xff\xff")

    def test_serialize_errors(self):
        with self.assertRaises(ProtocolError):
            Close(999, "").serialize()