"""
Tests for key ratcheting functionality.
"""

import pytest

from openmls_dave import KeyRatchet, KeyRatchetManager, SenderKeyRatchet


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


class TestKeyRatchet:
    """Tests for KeyRatchet class."""

    def test_create_ratchet(self, base_key):
        """Test creating a key ratchet."""
        ratchet = KeyRatchet(base_key)

        assert ratchet.generation == 0

    def test_create_ratchet_short_key_fails(self):
        """Test that creating ratchet with short key fails."""
        with pytest.raises(Exception):
            KeyRatchet(b"short")

    def test_get_current_key(self, key_ratchet):
        """Test getting current key."""
        key = to_bytes(key_ratchet.current_key())

        assert isinstance(key, bytes)
        assert len(key) == 32

    def test_get_key_at_generation(self, key_ratchet):
        """Test getting key at specific generation."""
        key0 = to_bytes(key_ratchet.get_key(0))
        key1 = to_bytes(key_ratchet.get_key(1))
        key2 = to_bytes(key_ratchet.get_key(2))

        # All keys should be different
        assert key0 != key1
        assert key1 != key2
        assert key0 != key2

    def test_ratchet_forward(self, key_ratchet):
        """Test ratcheting forward."""
        key0 = to_bytes(key_ratchet.current_key())

        key1 = to_bytes(key_ratchet.ratchet())
        assert key_ratchet.generation == 1
        assert key1 != key0

        key2 = to_bytes(key_ratchet.ratchet())
        assert key_ratchet.generation == 2
        assert key2 != key1

    def test_ratchet_to_generation(self, key_ratchet):
        """Test ratcheting to specific generation."""
        key = to_bytes(key_ratchet.ratchet_to(10))

        assert key_ratchet.generation == 10
        assert key == to_bytes(key_ratchet.get_key(10))

    def test_ratchet_backwards_fails(self, key_ratchet):
        """Test that ratcheting backwards fails."""
        key_ratchet.ratchet_to(10)

        with pytest.raises(Exception):
            key_ratchet.ratchet_to(5)

    def test_ratchet_large_gap_fails(self, key_ratchet):
        """Test that large generation gap fails."""
        with pytest.raises(Exception):
            key_ratchet.ratchet_to(1000)  # Gap > MAX_GENERATION_GAP

    def test_key_consistency(self, key_ratchet):
        """Test that same generation always produces same key."""
        key1 = to_bytes(key_ratchet.get_key(5))
        key2 = to_bytes(key_ratchet.get_key(5))

        assert key1 == key2

    def test_cache_truncation(self, base_key):
        """Test truncating the key cache."""
        ratchet = KeyRatchet(base_key, max_cached=16)

        # Generate several keys
        for i in range(20):
            ratchet.get_key(i)

        # Truncate old keys
        ratchet.truncate_cache(10)

        # Old keys should still be derivable (just not cached)
        key5 = to_bytes(ratchet.get_key(5))
        assert len(key5) == 32

    def test_clear_cache(self, key_ratchet):
        """Test clearing the key cache."""
        # Generate some keys
        for i in range(5):
            key_ratchet.get_key(i)

        key_ratchet.clear_cache()

        # Should still work after clearing
        key = to_bytes(key_ratchet.get_key(0))
        assert len(key) == 32


class TestSenderKeyRatchet:
    """Tests for SenderKeyRatchet class."""

    def test_create_sender_ratchet(self, base_key):
        """Test creating a sender key ratchet."""
        ratchet = SenderKeyRatchet(sender_id=12345, base_key=base_key)

        assert ratchet.sender_id == 12345
        assert ratchet.generation == 0

    def test_sender_ratchet_key(self, base_key):
        """Test getting key from sender ratchet."""
        ratchet = SenderKeyRatchet(sender_id=12345, base_key=base_key)
        key = to_bytes(ratchet.get_key(0))

        assert len(key) == 32

    def test_sender_ratchet_forward(self, base_key):
        """Test ratcheting sender key forward."""
        ratchet = SenderKeyRatchet(sender_id=12345, base_key=base_key)

        key0 = to_bytes(ratchet.get_key(0))
        key1 = to_bytes(ratchet.ratchet())

        assert key0 != key1
        assert ratchet.generation == 1


class TestKeyRatchetManager:
    """Tests for KeyRatchetManager class."""

    def test_create_manager(self, base_key):
        """Test creating a key ratchet manager."""
        manager = KeyRatchetManager(base_key)
        # Should not raise

    def test_get_or_create_ratchet(self, base_key):
        """Test getting or creating a ratchet."""
        manager = KeyRatchetManager(base_key)

        manager.get_or_create_ratchet(sender_id=1)
        manager.get_or_create_ratchet(sender_id=2)

        # Should not raise on duplicate
        manager.get_or_create_ratchet(sender_id=1)

    def test_get_key_for_sender(self, base_key):
        """Test getting key for specific sender."""
        manager = KeyRatchetManager(base_key)

        key1_gen0 = to_bytes(manager.get_key(sender_id=1, generation=0))
        key2_gen0 = to_bytes(manager.get_key(sender_id=2, generation=0))

        # Different senders should have different keys
        assert key1_gen0 != key2_gen0

    def test_get_key_different_generations(self, base_key):
        """Test getting keys at different generations."""
        manager = KeyRatchetManager(base_key)

        key_gen0 = to_bytes(manager.get_key(sender_id=1, generation=0))
        key_gen1 = to_bytes(manager.get_key(sender_id=1, generation=1))
        key_gen2 = to_bytes(manager.get_key(sender_id=1, generation=2))

        assert key_gen0 != key_gen1
        assert key_gen1 != key_gen2

    def test_remove_sender(self, base_key):
        """Test removing a sender."""
        manager = KeyRatchetManager(base_key)

        manager.get_key(sender_id=1, generation=0)
        manager.remove_sender(1)

        # Getting key again should work (recreates ratchet)
        key = to_bytes(manager.get_key(sender_id=1, generation=0))
        assert len(key) == 32

    def test_clear_all(self, base_key):
        """Test clearing all ratchets."""
        manager = KeyRatchetManager(base_key)

        manager.get_key(sender_id=1, generation=0)
        manager.get_key(sender_id=2, generation=0)

        manager.clear()

        # Should still work after clearing
        key = to_bytes(manager.get_key(sender_id=1, generation=0))
        assert len(key) == 32
