"""NTCP2 wire-format frame encryption.

Implements the NTCP2 wire format where each frame uses TWO AEAD
encrypt operations with the CipherState:

1. The frame length (2 bytes, big-endian) is encrypted first,
   producing 18 bytes (2 plaintext + 16 AEAD tag).
2. The frame payload (type + length + data) is encrypted next,
   producing N + 16 bytes (N plaintext + 16 AEAD tag).

Each operation consumes one nonce from the CipherState, so each
frame advances the nonce counter by 2.
"""

import struct

from i2p_crypto.noise import CipherState
from i2p_transport.ntcp2 import NTCP2Frame


class NTCP2WireCodec:
    """Encrypts and decrypts NTCP2 frames using the two-operation wire format."""

    def encrypt_frame(self, cipher: CipherState, frame: NTCP2Frame) -> bytes:
        """Encrypt an NTCP2Frame into wire-format bytes.

        Args:
            cipher: CipherState with key set (post-handshake).
            frame: The frame to encrypt.

        Returns:
            Wire bytes: encrypted_length (18) + encrypted_payload (N + 16).
        """
        frame_bytes = frame.to_bytes()

        # Operation 1: encrypt the 2-byte frame length
        length_plaintext = struct.pack("!H", len(frame_bytes))
        encrypted_length = cipher.encrypt_with_ad(b"", length_plaintext)

        # Operation 2: encrypt the frame bytes
        encrypted_payload = cipher.encrypt_with_ad(b"", frame_bytes)

        return encrypted_length + encrypted_payload

    def decrypt_frame_length(self, cipher: CipherState, encrypted_length: bytes) -> int:
        """Decrypt the 18-byte encrypted length field.

        Args:
            cipher: CipherState with matching key/nonce.
            encrypted_length: 18 bytes (2 plaintext + 16 AEAD tag).

        Returns:
            The frame length as an integer.
        """
        length_bytes = cipher.decrypt_with_ad(b"", encrypted_length)
        return struct.unpack("!H", length_bytes)[0]

    def decrypt_frame_payload(self, cipher: CipherState, encrypted_payload: bytes) -> NTCP2Frame:
        """Decrypt the encrypted payload and parse it as an NTCP2Frame.

        Args:
            cipher: CipherState with matching key/nonce (after length decrypt).
            encrypted_payload: The encrypted frame bytes (N + 16 bytes).

        Returns:
            The decrypted NTCP2Frame.
        """
        frame_bytes = cipher.decrypt_with_ad(b"", encrypted_payload)
        return NTCP2Frame.from_bytes(frame_bytes)

    def encrypt_and_get_wire_bytes(self, cipher: CipherState, frame: NTCP2Frame) -> bytes:
        """Convenience method: same as encrypt_frame.

        Args:
            cipher: CipherState with key set.
            frame: The frame to encrypt.

        Returns:
            Full wire bytes for transmission.
        """
        return self.encrypt_frame(cipher, frame)