"""Tests for the Double Ratchet implementation."""

import os
import pytest
from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey

from meshbay_common.ratchet import (
    RatchetState,
    MessageHeader,
    ChatMessage,
    encrypt_chat_message,
    decrypt_chat_message,
)


@pytest.fixture
def shared_secret():
    return os.urandom(32)


@pytest.fixture
def bob_key():
    return X25519PrivateKey.generate()


@pytest.fixture
def alice_bob(shared_secret, bob_key):
    from cryptography.hazmat.primitives import serialization
    bob_pub_raw = bob_key.public_key().public_bytes(
        serialization.Encoding.Raw, serialization.PublicFormat.Raw)
    alice = RatchetState.init_sender(shared_secret, bob_pub_raw)
    bob   = RatchetState.init_receiver(shared_secret, bob_key)
    return alice, bob


def test_basic_send_receive(alice_bob):
    alice, bob = alice_bob
    header, ct = alice.encrypt(b"Hello Bob!")
    plaintext = bob.decrypt(header, ct)
    assert plaintext == b"Hello Bob!"


def test_multiple_messages(alice_bob):
    alice, bob = alice_bob
    messages = [b"msg 1", b"msg 2", b"msg 3", b"msg 4", b"msg 5"]
    ciphertexts = [alice.encrypt(m) for m in messages]
    for i, (h, ct) in enumerate(ciphertexts):
        assert bob.decrypt(h, ct) == messages[i]


def test_bidirectional(alice_bob):
    """Both sides can send and receive."""
    alice, bob = alice_bob

    # Alice → Bob
    h, ct = alice.encrypt(b"Hi from Alice")
    assert bob.decrypt(h, ct) == b"Hi from Alice"

    # Bob → Alice (triggers DH ratchet on Alice)
    h, ct = bob.encrypt(b"Hi from Bob")
    assert alice.decrypt(h, ct) == b"Hi from Bob"

    # Alice → Bob again (new chain)
    h, ct = alice.encrypt(b"Second message from Alice")
    assert bob.decrypt(h, ct) == b"Second message from Alice"


def test_forward_secrecy(alice_bob):
    """Encrypting msg N+1 erases the key for msg N."""
    alice, bob = alice_bob
    h1, ct1 = alice.encrypt(b"message 1")
    h2, ct2 = alice.encrypt(b"message 2")

    assert bob.decrypt(h1, ct1) == b"message 1"
    assert bob.decrypt(h2, ct2) == b"message 2"

    # Keys are consumed — cannot replay
    with pytest.raises(Exception):
        bob.decrypt(h1, ct1)


def test_out_of_order_delivery(alice_bob):
    """Messages arriving out of order should still decrypt correctly."""
    alice, bob = alice_bob
    h1, ct1 = alice.encrypt(b"first")
    h2, ct2 = alice.encrypt(b"second")
    h3, ct3 = alice.encrypt(b"third")

    # Deliver in reverse order
    assert bob.decrypt(h3, ct3) == b"third"
    assert bob.decrypt(h2, ct2) == b"second"
    assert bob.decrypt(h1, ct1) == b"first"


def test_associated_data(alice_bob):
    alice, bob = alice_bob
    ad = b"sender:alice;group:test-group"
    h, ct = alice.encrypt(b"secret", associated_data=ad)
    assert bob.decrypt(h, ct, associated_data=ad) == b"secret"

    # Wrong AD fails authentication
    with pytest.raises(Exception):
        bob.decrypt(h, ct, associated_data=b"wrong-ad")


def test_header_encode_decode():
    from cryptography.hazmat.primitives import serialization
    sk = X25519PrivateKey.generate()
    pub = sk.public_key().public_bytes(
        serialization.Encoding.Raw, serialization.PublicFormat.Raw)
    h = MessageHeader(dh_pub=pub, prev_chain_n=5, msg_num=12)
    decoded = MessageHeader.decode(h.encode())
    assert decoded.dh_pub == pub
    assert decoded.prev_chain_n == 5
    assert decoded.msg_num == 12


def test_many_messages_stress(alice_bob):
    """100 messages without DH ratchet — verifies chain key stability."""
    alice, bob = alice_bob
    for i in range(100):
        h, ct = alice.encrypt(f"message {i}".encode())
        assert bob.decrypt(h, ct) == f"message {i}".encode()


def test_encrypt_decrypt_chat_message(alice_bob, shared_secret, bob_key):
    alice, bob = alice_bob
    msg = encrypt_chat_message(alice, "alice_user_id", "Hello group!")
    assert msg.sender_id == "alice_user_id"
    assert msg.msg_id != ""

    plaintext = decrypt_chat_message(bob, msg)
    assert plaintext == b"Hello group!"


def test_chat_message_serialization(alice_bob):
    alice, bob = alice_bob
    msg = encrypt_chat_message(alice, "alice", "Serialize me")
    d = msg.to_dict()
    restored = ChatMessage.from_dict(d)
    assert decrypt_chat_message(bob, restored) == b"Serialize me"


def test_break_in_recovery(shared_secret, bob_key):
    """
    Compromise of state at message N does not reveal keys for messages > N.
    After a DH ratchet step, new keys are independent of the compromised state.
    """
    from cryptography.hazmat.primitives import serialization
    bob_pub = bob_key.public_key().public_bytes(
        serialization.Encoding.Raw, serialization.PublicFormat.Raw)

    alice = RatchetState.init_sender(shared_secret, bob_pub)
    bob   = RatchetState.init_receiver(shared_secret, bob_key)

    # Exchange some messages
    h, ct = alice.encrypt(b"pre-compromise msg")
    bob.decrypt(h, ct)

    # Bob replies (triggers DH ratchet — new keys independent of above)
    h, ct = bob.encrypt(b"bob reply triggers ratchet")
    alice.decrypt(h, ct)

    # Now Alice's state has fresh keys
    h, ct = alice.encrypt(b"post-ratchet msg")
    plaintext = bob.decrypt(h, ct)
    assert plaintext == b"post-ratchet msg"