"""
Tests for SFrame encryption functionality.
"""

import pytest

from openmls_dave import SFrameContext, SFrameHeader, SFrameMultiContext


def to_bytes(data):
    """Convert data to bytes if it's a list."""
    if isinstance(data, list):
        return bytes(data)
    return data


class TestSFrameHeader:
    """Tests for SFrame header encoding/decoding."""

    def test_create_header(self):
        """Test creating an SFrame header."""
        header = SFrameHeader(key_id=1, counter=42)

        assert header.key_id == 1
        assert header.counter == 42

    def test_serialize_compact_header(self):
        """Test serializing a compact header (small values)."""
        header = SFrameHeader(key_id=3, counter=7)
        serialized = to_bytes(header.serialize())

        # Compact format should be 1 byte
        assert len(serialized) == 1

    def test_serialize_extended_header(self):
        """Test serializing an extended header (large values)."""
        header = SFrameHeader(key_id=256, counter=65536)
        serialized = to_bytes(header.serialize())

        # Extended format needs more bytes
        assert len(serialized) > 1

    def test_parse_compact_header(self):
        """Test parsing a compact header."""
        header = SFrameHeader(key_id=5, counter=10)
        serialized = to_bytes(header.serialize())

        parsed, consumed = SFrameHeader.parse(serialized)

        assert parsed.key_id == 5
        assert parsed.counter == 10
        assert consumed == len(serialized)

    def test_parse_extended_header(self):
        """Test parsing an extended header."""
        header = SFrameHeader(key_id=1000, counter=100000)
        serialized = to_bytes(header.serialize())

        parsed, consumed = SFrameHeader.parse(serialized)

        assert parsed.key_id == 1000
        assert parsed.counter == 100000

    def test_roundtrip_various_values(self):
        """Test header roundtrip with various values."""
        test_cases = [
            (0, 0),
            (1, 1),
            (7, 15),
            (8, 16),
            (255, 255),
            (256, 256),
            (65535, 65535),
            (0xFFFFFFFF, 0xFFFFFFFF),
        ]

        for key_id, counter in test_cases:
            header = SFrameHeader(key_id=key_id, counter=counter)
            serialized = to_bytes(header.serialize())
            parsed, _ = SFrameHeader.parse(serialized)

            assert parsed.key_id == key_id, f"key_id mismatch for ({key_id}, {counter})"
            assert parsed.counter == counter, f"counter mismatch for ({key_id}, {counter})"


class TestSFrameContext:
    """Tests for SFrame encryption context."""

    def test_create_context(self, base_key):
        """Test creating an SFrame context."""
        ctx = SFrameContext(to_bytes(base_key), key_id=1)

        assert ctx.key_id == 1
        assert ctx.current_counter == 0

    def test_create_context_short_key_fails(self):
        """Test that creating context with short key fails."""
        with pytest.raises(Exception):
            SFrameContext(b"short", key_id=1)

    def test_encrypt_frame(self, sframe_context, sample_audio_frame):
        """Test encrypting a frame."""
        ciphertext = to_bytes(sframe_context.encrypt(to_bytes(sample_audio_frame)))

        assert isinstance(ciphertext, bytes)
        assert len(ciphertext) > len(sample_audio_frame)  # Header + tag overhead
        assert ciphertext != sample_audio_frame

    def test_decrypt_frame(self, sframe_context, sample_audio_frame):
        """Test decrypting a frame."""
        sample = to_bytes(sample_audio_frame)
        ciphertext = to_bytes(sframe_context.encrypt(sample))
        plaintext, header = sframe_context.decrypt(ciphertext)
        plaintext = to_bytes(plaintext)

        assert plaintext == sample
        assert header.key_id == sframe_context.key_id

    def test_encrypt_increments_counter(self, sframe_context, sample_audio_frame):
        """Test that encryption increments counter."""
        sample = to_bytes(sample_audio_frame)
        initial_counter = sframe_context.current_counter

        sframe_context.encrypt(sample)
        assert sframe_context.current_counter == initial_counter + 1

        sframe_context.encrypt(sample)
        assert sframe_context.current_counter == initial_counter + 2

    def test_encrypt_with_specific_counter(self, sframe_context, sample_audio_frame):
        """Test encrypting with a specific counter value."""
        sample = to_bytes(sample_audio_frame)
        ciphertext = to_bytes(sframe_context.encrypt_with_counter(sample, counter=42))
        plaintext, header = sframe_context.decrypt(ciphertext)
        plaintext = to_bytes(plaintext)

        assert plaintext == sample
        assert header.counter == 42

    def test_decrypt_wrong_key_fails(self, base_key, sample_audio_frame):
        """Test that decryption with wrong key fails."""
        key1 = to_bytes(base_key)
        key2 = bytes(reversed(key1))
        sample = to_bytes(sample_audio_frame)

        ctx1 = SFrameContext(key1, key_id=1)
        ctx2 = SFrameContext(key2, key_id=1)

        ciphertext = to_bytes(ctx1.encrypt(sample))

        with pytest.raises(Exception):
            ctx2.decrypt(ciphertext)

    def test_decrypt_wrong_key_id_fails(self, base_key, sample_audio_frame):
        """Test that decryption with wrong key ID fails."""
        key = to_bytes(base_key)
        sample = to_bytes(sample_audio_frame)

        ctx1 = SFrameContext(key, key_id=1)
        ctx2 = SFrameContext(key, key_id=2)

        ciphertext = to_bytes(ctx1.encrypt(sample))

        with pytest.raises(Exception):
            ctx2.decrypt(ciphertext)

    def test_set_counter(self, sframe_context):
        """Test setting the counter."""
        sframe_context.set_counter(100)
        assert sframe_context.current_counter == 100

    def test_large_frame(self, sframe_context, sample_video_frame):
        """Test encrypting a large frame."""
        sample = to_bytes(sample_video_frame)
        ciphertext = to_bytes(sframe_context.encrypt(sample))
        plaintext, _ = sframe_context.decrypt(ciphertext)
        plaintext = to_bytes(plaintext)

        assert plaintext == sample

    def test_empty_frame(self, sframe_context):
        """Test encrypting an empty frame."""
        ciphertext = to_bytes(sframe_context.encrypt(b""))
        plaintext, _ = sframe_context.decrypt(ciphertext)
        plaintext = to_bytes(plaintext)

        assert plaintext == b""

    def test_context_repr(self, sframe_context):
        """Test context string representation."""
        repr_str = repr(sframe_context)

        assert "SFrameContext" in repr_str
        assert "key_id=" in repr_str


class TestSFrameMultiContext:
    """Tests for multi-sender SFrame context."""

    def test_create_multi_context(self):
        """Test creating a multi-sender context."""
        ctx = SFrameMultiContext(local_key_id=1)
        assert ctx.sender_count() == 0

    def test_add_sender(self, base_key):
        """Test adding a sender."""
        ctx = SFrameMultiContext(local_key_id=1)
        ctx.add_sender(1, to_bytes(base_key))

        assert ctx.has_sender(1)
        assert ctx.sender_count() == 1

    def test_remove_sender(self, base_key):
        """Test removing a sender."""
        key = to_bytes(base_key)
        ctx = SFrameMultiContext(local_key_id=1)
        ctx.add_sender(1, key)
        ctx.add_sender(2, key)

        ctx.remove_sender(1)

        assert not ctx.has_sender(1)
        assert ctx.has_sender(2)

    def test_encrypt_decrypt_local(self, base_key, sample_audio_frame):
        """Test encrypting and decrypting with local key."""
        key = to_bytes(base_key)
        sample = to_bytes(sample_audio_frame)

        ctx = SFrameMultiContext(local_key_id=1)
        ctx.add_sender(1, key)

        ciphertext = to_bytes(ctx.encrypt(sample))
        plaintext, header = ctx.decrypt(ciphertext)
        plaintext = to_bytes(plaintext)

        assert plaintext == sample
        assert header.key_id == 1

    def test_decrypt_from_multiple_senders(self, sample_audio_frame):
        """Test decrypting frames from multiple senders."""
        key1 = bytes(range(32))
        key2 = bytes(range(32, 64))

        # Sender 1's context
        sender1_ctx = SFrameContext(key1, key_id=1)

        # Sender 2's context
        sender2_ctx = SFrameContext(key2, key_id=2)

        # Receiver's multi-context
        receiver = SFrameMultiContext(local_key_id=99)
        receiver.add_sender(1, key1)
        receiver.add_sender(2, key2)

        # Encrypt from both senders
        ct1 = to_bytes(sender1_ctx.encrypt(b"from sender 1"))
        ct2 = to_bytes(sender2_ctx.encrypt(b"from sender 2"))

        # Receiver can decrypt both
        pt1, hdr1 = receiver.decrypt(ct1)
        pt2, hdr2 = receiver.decrypt(ct2)

        pt1 = to_bytes(pt1)
        pt2 = to_bytes(pt2)

        assert pt1 == b"from sender 1"
        assert hdr1.key_id == 1

        assert pt2 == b"from sender 2"
        assert hdr2.key_id == 2

    def test_decrypt_unknown_sender_fails(self, base_key, sample_audio_frame):
        """Test that decrypting from unknown sender fails."""
        key = to_bytes(base_key)
        sample = to_bytes(sample_audio_frame)

        sender = SFrameContext(key, key_id=1)
        receiver = SFrameMultiContext(local_key_id=99)
        # Don't add sender 1

        ciphertext = to_bytes(sender.encrypt(sample))

        with pytest.raises(Exception):
            receiver.decrypt(ciphertext)
