"""Unit tests for WebSocket handshake and frame behavior."""

from __future__ import annotations

import base64
import asyncio
import hashlib
import struct
import unittest

from src.errors import RdpProtocolError
from src.websocket_client import (
    FrameDecoder,
    WEBSOCKET_GUID,
    WebSocketClient,
    encodeFrame,
)


def serverFrame(payload: bytes, opcode: int = 1, final: bool = True) -> bytes:
    """Build one unmasked server frame for decoder tests."""
    first = (0x80 if final else 0) | opcode
    size = len(payload)
    if size < 126:
        return bytes((first, size)) + payload
    if size <= 65_535:
        return bytes((first, 126)) + struct.pack("!H", size) + payload
    return bytes((first, 127)) + struct.pack("!Q", size) + payload


class FrameDecoderTest(unittest.TestCase):
    """Verify incremental frame decoding and RFC validation."""

    def test_every_split_boundary(self) -> None:
        frame = serverFrame("hello 😀".encode())
        for offset in range(len(frame) + 1):
            decoder = FrameDecoder()
            first = decoder.feed(frame[:offset])
            second = decoder.feed(frame[offset:])
            complete = first + second
            self.assertEqual(complete[-1].payload, "hello 😀".encode())

    def test_extended_lengths(self) -> None:
        for size in (126, 65_536):
            payload = b"x" * size
            decoded = FrameDecoder(max_message_bytes=100_000).feed(
                serverFrame(payload)
            )
            self.assertEqual(decoded[-1].payload, payload)

    def test_fragmented_text_with_ping(self) -> None:
        decoder = FrameDecoder()
        frames = []
        frames.extend(decoder.feed(serverFrame(b"hel", 1, False)))
        frames.extend(decoder.feed(serverFrame(b"p", 9, True)))
        frames.extend(decoder.feed(serverFrame(b"lo", 0, True)))
        self.assertEqual(frames[-1].payload, b"hello")
        self.assertEqual(frames[-2].opcode, 9)

    def test_rejects_invalid_frames(self) -> None:
        cases = [
            bytes((0xC1, 0)),
            bytes((0x81, 0x80)) + b"mask",
            bytes((0x09, 0)),
            serverFrame(b"\xff", 1),
            serverFrame(b"binary", 2),
            serverFrame(b"continuation", 0),
        ]
        for frame in cases:
            with self.subTest(frame=frame[:2]):
                with self.assertRaises(RdpProtocolError):
                    FrameDecoder().feed(frame)

    def test_rejects_oversized_message(self) -> None:
        with self.assertRaises(RdpProtocolError):
            FrameDecoder(max_message_bytes=3).feed(serverFrame(b"four"))


class FrameEncoderTest(unittest.TestCase):
    """Verify mandatory client masking and encoded length forms."""

    def test_masks_payload(self) -> None:
        encoded = encodeFrame(b"hello", mask_key=b"abcd")
        self.assertEqual(encoded[:2], bytes((0x81, 0x85)))
        self.assertEqual(encoded[2:6], b"abcd")
        unmasked = bytes(
            value ^ b"abcd"[index % 4]
            for index, value in enumerate(encoded[6:])
        )
        self.assertEqual(unmasked, b"hello")

    def test_uses_extended_lengths(self) -> None:
        medium = encodeFrame(b"x" * 126, mask_key=b"abcd")
        large = encodeFrame(b"x" * 65_536, mask_key=b"abcd")
        self.assertEqual(medium[1] & 0x7f, 126)
        self.assertEqual(large[1] & 0x7f, 127)


class HandshakeTest(unittest.TestCase):
    """Verify strict HTTP Upgrade response validation."""

    def test_accepts_valid_response(self) -> None:
        key = "dGhlIHNhbXBsZSBub25jZQ=="
        accept = base64.b64encode(
            hashlib.sha1((key + WEBSOCKET_GUID).encode()).digest()
        ).decode()
        response = (
            "HTTP/1.1 101 Switching Protocols\r\n"
            "Upgrade: WebSocket\r\n"
            "Connection: keep-alive, Upgrade\r\n"
            f"Sec-WebSocket-Accept: {accept}\r\n\r\n"
        ).encode()
        WebSocketClient._validateHandshake(response, key)

    def test_rejects_incorrect_accept(self) -> None:
        response = (
            b"HTTP/1.1 101 Switching Protocols\r\n"
            b"Upgrade: websocket\r\nConnection: Upgrade\r\n"
            b"Sec-WebSocket-Accept: wrong\r\n\r\n"
        )
        with self.assertRaises(RdpProtocolError):
            WebSocketClient._validateHandshake(response, "key")


class ClientReceiveTest(unittest.IsolatedAsyncioTestCase):
    """Verify the stream client accumulates partial frame reads."""

    async def test_partial_frame_reads_do_not_pop_empty_queue(self) -> None:
        frame = serverFrame(b"split message")

        class SplitReader:
            """Return a frame in pieces smaller than its header and payload."""

            def __init__(self) -> None:
                self.parts = [frame[:1], frame[1:3], frame[3:]]

            async def read(self, size: int) -> bytes:
                """Return the next configured stream fragment."""
                await asyncio.sleep(0)
                return self.parts.pop(0) if self.parts else b""

        client = WebSocketClient("127.0.0.1", 6000)
        client._reader = SplitReader()
        self.assertEqual(await client.receiveText(), "split message")


if __name__ == "__main__":
    unittest.main()