"""
Tests for the Sender Keys group messaging protocol.

Covers: key creation, distribution, encrypt/decrypt, multi-member groups,
out-of-order delivery, serialization, and key rotation on member removal.
"""

import pytest

from meshbay_common.senderkeys import (
    SenderKeyRecord,
    SenderKeyDistribution,
    SenderKeyMessage,
    GroupSenderKeyStore,
    encrypt_message,
    decrypt_message,
)


def test_basic_encrypt_decrypt():
    """Alice encrypts, Bob decrypts using Alice's distributed sender key."""
    alice_rec = SenderKeyRecord.create("alice")
    alice_dist = alice_rec.distribution()

    bob_store = GroupSenderKeyStore("group-1")
    bob_store.add_sender(alice_dist)

    msg, alice_rec = encrypt_message(alice_rec, b"hello group")
    plaintext = decrypt_message(bob_store, msg)
    assert plaintext == b"hello group"


def test_multiple_messages_sequential():
    """Multiple messages from the same sender decrypt in order."""
    alice_rec = SenderKeyRecord.create("alice")
    store = GroupSenderKeyStore("group-1")
    store.add_sender(alice_rec.distribution())

    for i in range(5):
        msg, alice_rec = encrypt_message(alice_rec, f"message {i}".encode())
        pt = decrypt_message(store, msg)
        assert pt == f"message {i}".encode()


def test_multi_member_group():
    """Three members: Alice sends, Bob and Carol both decrypt."""
    alice_rec = SenderKeyRecord.create("alice")
    alice_dist = alice_rec.distribution()

    bob_store = GroupSenderKeyStore("group-1")
    bob_store.add_sender(alice_dist)

    carol_store = GroupSenderKeyStore("group-1")
    carol_store.add_sender(alice_dist)

    msg, alice_rec = encrypt_message(alice_rec, b"broadcast")

    assert decrypt_message(bob_store, msg) == b"broadcast"
    assert decrypt_message(carol_store, msg) == b"broadcast"


def test_bidirectional_chat():
    """Alice and Bob both send and receive."""
    alice_rec = SenderKeyRecord.create("alice")
    bob_rec = SenderKeyRecord.create("bob")

    alice_store = GroupSenderKeyStore("group-1")
    alice_store.add_sender(bob_rec.distribution())

    bob_store = GroupSenderKeyStore("group-1")
    bob_store.add_sender(alice_rec.distribution())

    msg1, alice_rec = encrypt_message(alice_rec, b"hi bob")
    assert decrypt_message(bob_store, msg1) == b"hi bob"

    msg2, bob_rec = encrypt_message(bob_rec, b"hi alice")
    assert decrypt_message(alice_store, msg2) == b"hi alice"


def test_out_of_order_delivery():
    """Messages delivered out of order are decrypted correctly (up to MAX_SKIP)."""
    alice_rec = SenderKeyRecord.create("alice")
    store = GroupSenderKeyStore("group-1")
    store.add_sender(alice_rec.distribution())

    msg0, alice_rec = encrypt_message(alice_rec, b"msg 0")
    msg1, alice_rec = encrypt_message(alice_rec, b"msg 1")
    msg2, alice_rec = encrypt_message(alice_rec, b"msg 2")

    # Deliver out of order: 2, 0, 1
    assert decrypt_message(store, msg2) == b"msg 2"
    assert decrypt_message(store, msg0) == b"msg 0"
    assert decrypt_message(store, msg1) == b"msg 1"


def test_replay_rejected():
    """A message decrypted twice raises an error (replay protection)."""
    alice_rec = SenderKeyRecord.create("alice")
    store = GroupSenderKeyStore("group-1")
    store.add_sender(alice_rec.distribution())

    msg, alice_rec = encrypt_message(alice_rec, b"once only")
    decrypt_message(store, msg)

    with pytest.raises(ValueError, match="already consumed"):
        decrypt_message(store, msg)


def test_unknown_sender_rejected():
    """Message from an unknown sender raises ValueError."""
    alice_rec = SenderKeyRecord.create("alice")
    store = GroupSenderKeyStore("group-1")

    msg, _ = encrypt_message(alice_rec, b"who am i")
    with pytest.raises(ValueError, match="Unknown sender"):
        decrypt_message(store, msg)


def test_non_member_cannot_decrypt():
    """Eve (not in group) cannot decrypt Alice's messages."""
    alice_rec = SenderKeyRecord.create("alice")
    eve_store = GroupSenderKeyStore("group-1")

    msg, _ = encrypt_message(alice_rec, b"secret")
    with pytest.raises(ValueError, match="Unknown sender"):
        decrypt_message(eve_store, msg)


def test_key_rotation_on_member_removal():
    """After rotation, old chain keys cannot decrypt new messages."""
    alice_rec = SenderKeyRecord.create("alice")
    old_dist = alice_rec.distribution()

    # Eve had Alice's old key
    eve_store = GroupSenderKeyStore("group-1")
    eve_store.add_sender(old_dist)

    # Alice rotates (member removed from group)
    alice_rec = alice_rec.rotate()
    new_dist = alice_rec.distribution()

    # Bob gets the new distribution
    bob_store = GroupSenderKeyStore("group-1")
    bob_store.add_sender(new_dist)

    msg, alice_rec = encrypt_message(alice_rec, b"post-rotation")
    assert decrypt_message(bob_store, msg) == b"post-rotation"

    # Eve cannot decrypt with old key
    with pytest.raises(Exception):
        decrypt_message(eve_store, msg)


def test_distribution_serialization():
    """SenderKeyDistribution round-trips through serialize/deserialize."""
    rec = SenderKeyRecord.create("alice")
    dist = rec.distribution()
    data = dist.serialize()
    recovered = SenderKeyDistribution.deserialize(data)

    assert recovered.sender_id == dist.sender_id
    assert recovered.chain_key == dist.chain_key
    assert recovered.iteration == dist.iteration
    assert recovered.signing_pk == dist.signing_pk


def test_message_serialization():
    """SenderKeyMessage round-trips through serialize/deserialize."""
    rec = SenderKeyRecord.create("alice")
    store = GroupSenderKeyStore("group-1")
    store.add_sender(rec.distribution())

    msg, _ = encrypt_message(rec, b"serialize me")
    data = msg.serialize()
    recovered = SenderKeyMessage.deserialize(data)

    assert recovered.sender_id == msg.sender_id
    assert recovered.iteration == msg.iteration
    assert recovered.ciphertext == msg.ciphertext
    assert recovered.nonce == msg.nonce
    assert recovered.signature == msg.signature

    # Deserialized message still decrypts
    pt = decrypt_message(store, recovered)
    assert pt == b"serialize me"


def test_tampered_ciphertext_rejected():
    """Modifying the ciphertext makes signature verification fail."""
    alice_rec = SenderKeyRecord.create("alice")
    store = GroupSenderKeyStore("group-1")
    store.add_sender(alice_rec.distribution())

    msg, _ = encrypt_message(alice_rec, b"authentic")
    msg.ciphertext = bytes([b ^ 0xff for b in msg.ciphertext])

    with pytest.raises(Exception):
        decrypt_message(store, msg)


def test_store_sender_count():
    """GroupSenderKeyStore tracks sender count correctly."""
    store = GroupSenderKeyStore("group-1")
    assert store.sender_count == 0

    store.add_sender(SenderKeyRecord.create("alice").distribution())
    store.add_sender(SenderKeyRecord.create("bob").distribution())
    assert store.sender_count == 2

    store.remove_sender("alice")
    assert store.sender_count == 1