"""Tests for Noise_XX handshake and encrypted framing."""

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

from a2a_p2p.crypto import (
    NoiseInitiator,
    NoiseResponder,
    NoiseSession,
    frame_decrypt,
    frame_decrypt_length,
    frame_encrypt,
)


class TestNoiseXX:
    def _do_handshake(self):
        """Perform a full Noise_XX handshake between two parties."""
        init_static = X25519PrivateKey.generate()
        resp_static = X25519PrivateKey.generate()

        initiator = NoiseInitiator(init_static)
        responder = NoiseResponder(resp_static)

        # Message 1: initiator → responder
        msg1 = initiator.write_message_1()

        # Message 2: responder → initiator
        msg2 = responder.read_message_1(msg1)

        # Message 3: initiator → responder
        msg3 = initiator.read_message_2(msg2)
        responder.read_message_3(msg3)

        init_session = initiator.finalize()
        resp_session = responder.finalize()

        return init_session, resp_session

    def test_handshake_completes(self):
        init_session, resp_session = self._do_handshake()
        assert isinstance(init_session, NoiseSession)
        assert isinstance(resp_session, NoiseSession)

    def test_mutual_key_agreement(self):
        """Both sides should derive matching remote static keys."""
        init_static = X25519PrivateKey.generate()
        resp_static = X25519PrivateKey.generate()

        from cryptography.hazmat.primitives.serialization import Encoding, PublicFormat

        init_pub = init_static.public_key().public_bytes(Encoding.Raw, PublicFormat.Raw)
        resp_pub = resp_static.public_key().public_bytes(Encoding.Raw, PublicFormat.Raw)

        initiator = NoiseInitiator(init_static)
        responder = NoiseResponder(resp_static)

        msg1 = initiator.write_message_1()
        msg2 = responder.read_message_1(msg1)
        msg3 = initiator.read_message_2(msg2)
        responder.read_message_3(msg3)

        init_session = initiator.finalize()
        resp_session = responder.finalize()

        # Initiator should see responder's static key
        assert init_session.remote_static_pub == resp_pub
        # Responder should see initiator's static key
        assert resp_session.remote_static_pub == init_pub

    def test_encrypted_communication(self):
        """After handshake, messages should encrypt/decrypt correctly."""
        init_session, resp_session = self._do_handshake()

        # Initiator → Responder
        plaintext = b"Hello from initiator!"
        ct = init_session.encrypt(plaintext)
        assert ct != plaintext
        pt = resp_session.decrypt(ct)
        assert pt == plaintext

        # Responder → Initiator
        plaintext2 = b"Hello from responder!"
        ct2 = resp_session.encrypt(plaintext2)
        pt2 = init_session.decrypt(ct2)
        assert pt2 == plaintext2

    def test_multiple_messages(self):
        """Multiple sequential messages should work (nonce increments)."""
        init_session, resp_session = self._do_handshake()

        for i in range(10):
            msg = f"Message {i}".encode()
            ct = init_session.encrypt(msg)
            pt = resp_session.decrypt(ct)
            assert pt == msg

    def test_tampered_ciphertext_fails(self):
        init_session, resp_session = self._do_handshake()

        ct = init_session.encrypt(b"secret")
        # Tamper
        tampered = bytearray(ct)
        tampered[-1] ^= 0xFF

        import pytest
        from a2a_p2p.errors import CryptoError

        with pytest.raises(CryptoError):
            resp_session.decrypt(bytes(tampered))


class TestFraming:
    def test_frame_roundtrip(self):
        init_static = X25519PrivateKey.generate()
        resp_static = X25519PrivateKey.generate()

        initiator = NoiseInitiator(init_static)
        responder = NoiseResponder(resp_static)

        msg1 = initiator.write_message_1()
        msg2 = responder.read_message_1(msg1)
        msg3 = initiator.read_message_2(msg2)
        responder.read_message_3(msg3)

        init_session = initiator.finalize()
        resp_session = responder.finalize()

        plaintext = b"framed message test"
        framed = frame_encrypt(init_session, plaintext)

        # Parse length
        length = frame_decrypt_length(framed[:4])
        ct = framed[4:]
        assert len(ct) == length

        pt = frame_decrypt(resp_session, ct)
        assert pt == plaintext
