"""
High-level Pythonic API for AEGIS authenticated encryption.
"""

import os

from pyaegis._aegis_ffi import ffi, lib


class AegisError(Exception):
    """Base exception for AEGIS operations."""

    pass


class DecryptionError(AegisError):
    """Raised when decryption or authentication fails."""

    pass


# Initialize the library on import (runtime CPU detection)
lib.aegis_init()


class _AEGISBase:
    """Base class for AEGIS cipher implementations with shared functionality."""

    KEY_SIZE: int
    NONCE_SIZE: int
    TAG_SIZE_MIN: int = 16
    TAG_SIZE_MAX: int = 32

    # Subclasses should define these as references to C library functions
    _encrypt_func = None
    _decrypt_func = None
    _encrypt_detached_func = None
    _decrypt_detached_func = None
    _stream_func = None

    def __init__(self, tag_size: int = 32) -> None:
        """
        Initialize an AEGIS cipher.

        Args:
            tag_size: Size of the authentication tag in bytes (16 or 32).
                     32 bytes is recommended for maximum security.

        Raises:
            ValueError: If tag_size is not 16 or 32.
        """
        if tag_size not in (16, 32):
            raise ValueError(f"tag_size must be 16 or 32, got {tag_size}")
        self.tag_size = tag_size

    @classmethod
    def random_key(cls) -> bytes:
        """Generate a random key suitable for this cipher."""
        return os.urandom(cls.KEY_SIZE)

    @classmethod
    def random_nonce(cls) -> bytes:
        """Generate a random nonce suitable for this cipher."""
        return os.urandom(cls.NONCE_SIZE)

    def _check_key(self, key: bytes) -> None:
        """Validate key length."""
        if len(key) != self.KEY_SIZE:
            raise ValueError(f"Key must be {self.KEY_SIZE} bytes, got {len(key)}")

    def _check_nonce(self, nonce: bytes) -> None:
        """Validate nonce length."""
        if len(nonce) != self.NONCE_SIZE:
            raise ValueError(f"Nonce must be {self.NONCE_SIZE} bytes, got {len(nonce)}")

    def encrypt(
        self,
        key: bytes,
        nonce: bytes,
        plaintext: bytes,
        associated_data: bytes | None = None,
        into: bytearray | None = None,
    ) -> bytes:
        """
        Encrypt and authenticate a message.

        Args:
            key: Encryption key (size depends on cipher variant)
            nonce: Nonce (must be unique for each message with the same key)
            plaintext: Message to encrypt
            associated_data: Optional additional authenticated data (not encrypted)
            into: Optional pre-allocated bytearray for output (must be len(plaintext) + tag_size)

        Returns:
            Ciphertext with authentication tag appended

        Raises:
            ValueError: If key or nonce has incorrect length, or if into has wrong size
            AegisError: If encryption fails
        """
        self._check_key(key)
        self._check_nonce(nonce)

        plaintext_len = len(plaintext)
        output_len = plaintext_len + self.tag_size

        # Use pre-allocated buffer if provided, otherwise allocate
        if into is not None:
            if len(into) != output_len:
                raise ValueError(f"Output buffer must be {output_len} bytes, got {len(into)}")
            output_buf = ffi.from_buffer(into)
        else:
            into = bytearray(output_len)
            output_buf = ffi.from_buffer(into)

        if associated_data:
            ad_len = len(associated_data)
            ad_ptr = associated_data if ad_len > 0 else ffi.NULL
        else:
            ad_len = 0
            ad_ptr = ffi.NULL

        # Multi-lane variants use detached mode internally
        if self._encrypt_func is None:
            assert self._encrypt_detached_func is not None
            ciphertext_buf = output_buf
            tag_buf = ffi.cast("uint8_t*", output_buf) + plaintext_len

            result = self._encrypt_detached_func(
                ciphertext_buf,
                tag_buf,
                self.tag_size,
                plaintext,
                plaintext_len,
                ad_ptr,
                ad_len,
                nonce,
                key,
            )

            if result != 0:
                raise AegisError("Encryption failed")

            return bytes(into)

        # Standard variants use combined mode
        result = self._encrypt_func(
            output_buf,
            self.tag_size,
            plaintext,
            plaintext_len,
            ad_ptr,
            ad_len,
            nonce,
            key,
        )

        if result != 0:
            raise AegisError("Encryption failed")

        return bytes(into)

    def decrypt(
        self,
        key: bytes,
        nonce: bytes,
        ciphertext: bytes,
        associated_data: bytes | None = None,
        into: bytearray | None = None,
    ) -> bytes:
        """
        Decrypt and verify an authenticated message.

        Args:
            key: Encryption key (size depends on cipher variant)
            nonce: Nonce used during encryption
            ciphertext: Encrypted message with authentication tag
            associated_data: Optional additional authenticated data
            into: Optional pre-allocated bytearray for output (must be len(ciphertext) - tag_size)

        Returns:
            Decrypted plaintext

        Raises:
            ValueError: If key or nonce has incorrect length, or if into has wrong size
            DecryptionError: If authentication fails
        """
        self._check_key(key)
        self._check_nonce(nonce)

        ciphertext_len = len(ciphertext)
        if ciphertext_len < self.tag_size:
            raise DecryptionError("Ciphertext too short")

        plaintext_len = ciphertext_len - self.tag_size

        # Use pre-allocated buffer if provided, otherwise allocate
        if into is not None:
            if len(into) != plaintext_len:
                raise ValueError(f"Output buffer must be {plaintext_len} bytes, got {len(into)}")
            plaintext_buf = ffi.from_buffer(into)
        else:
            into = bytearray(plaintext_len)
            plaintext_buf = ffi.from_buffer(into)

        if associated_data:
            ad_len = len(associated_data)
            ad_ptr = associated_data if ad_len > 0 else ffi.NULL
        else:
            ad_len = 0
            ad_ptr = ffi.NULL

        # Multi-lane variants use detached mode internally
        if self._decrypt_func is None:
            assert self._decrypt_detached_func is not None
            ct_ptr = ffi.cast("const uint8_t*", ffi.from_buffer(ciphertext))
            tag_ptr = ct_ptr + plaintext_len

            result = self._decrypt_detached_func(
                plaintext_buf,
                ct_ptr,
                plaintext_len,
                tag_ptr,
                self.tag_size,
                ad_ptr,
                ad_len,
                nonce,
                key,
            )
        else:
            # Standard variants use combined mode
            result = self._decrypt_func(
                plaintext_buf,
                ciphertext,
                ciphertext_len,
                self.tag_size,
                ad_ptr,
                ad_len,
                nonce,
                key,
            )

        if result != 0:
            raise DecryptionError("Authentication failed")

        return bytes(into)


class _AEGISWithDetached(_AEGISBase):
    """Base class for AEGIS variants that support detached mode explicitly."""

    def encrypt_detached(
        self,
        key: bytes,
        nonce: bytes,
        plaintext: bytes,
        associated_data: bytes | None = None,
        ciphertext_into: bytearray | None = None,
    ) -> tuple[bytes, bytes]:
        """
        Encrypt and authenticate a message, returning ciphertext and tag separately.

        Args:
            key: Encryption key
            nonce: Nonce
            plaintext: Message to encrypt
            associated_data: Optional additional authenticated data
            ciphertext_into: Optional pre-allocated bytearray for ciphertext (must be len(plaintext))

        Returns:
            Tuple of (ciphertext, authentication_tag)

        Raises:
            ValueError: If key or nonce has incorrect length, or if ciphertext_into has wrong size
            AegisError: If encryption fails
        """
        self._check_key(key)
        self._check_nonce(nonce)

        if associated_data is None:
            associated_data = b""

        plaintext_len = len(plaintext)

        # Use pre-allocated buffer for ciphertext if provided
        if ciphertext_into is not None:
            if len(ciphertext_into) != plaintext_len:
                raise ValueError(
                    f"Ciphertext buffer must be {plaintext_len} bytes, got {len(ciphertext_into)}"
                )
            ciphertext_buf = ffi.from_buffer(ciphertext_into)
        else:
            ciphertext_into = bytearray(plaintext_len)
            ciphertext_buf = ffi.from_buffer(ciphertext_into)

        # Tag is small (16-32 bytes), so just allocate it
        tag_into = bytearray(self.tag_size)
        tag_buf = ffi.from_buffer(tag_into)

        assert self._encrypt_detached_func is not None

        result = self._encrypt_detached_func(
            ciphertext_buf,
            tag_buf,
            self.tag_size,
            plaintext,
            plaintext_len,
            associated_data if len(associated_data) > 0 else ffi.NULL,
            len(associated_data),
            nonce,
            key,
        )

        if result != 0:
            raise AegisError("Encryption failed")

        return (bytes(ciphertext_into), bytes(tag_into))

    def decrypt_detached(
        self,
        key: bytes,
        nonce: bytes,
        ciphertext: bytes,
        tag: bytes,
        associated_data: bytes | None = None,
        into: bytearray | None = None,
    ) -> bytes:
        """
        Decrypt and verify a message with detached authentication tag.

        Args:
            key: Encryption key
            nonce: Nonce
            ciphertext: Encrypted message
            tag: Authentication tag
            associated_data: Optional additional authenticated data
            into: Optional pre-allocated bytearray for plaintext (must be len(ciphertext))

        Returns:
            Decrypted plaintext

        Raises:
            ValueError: If key or nonce has incorrect length, or if into has wrong size
            DecryptionError: If authentication fails
        """
        self._check_key(key)
        self._check_nonce(nonce)

        if associated_data is None:
            associated_data = b""

        ciphertext_len = len(ciphertext)

        # Use pre-allocated buffer if provided
        if into is not None:
            if len(into) != ciphertext_len:
                raise ValueError(f"Output buffer must be {ciphertext_len} bytes, got {len(into)}")
            plaintext_buf = ffi.from_buffer(into)
        else:
            into = bytearray(ciphertext_len)
            plaintext_buf = ffi.from_buffer(into)

        assert self._decrypt_detached_func is not None

        result = self._decrypt_detached_func(
            plaintext_buf,
            ciphertext,
            ciphertext_len,
            tag,
            len(tag),
            associated_data if len(associated_data) > 0 else ffi.NULL,
            len(associated_data),
            nonce,
            key,
        )

        if result != 0:
            raise DecryptionError("Authentication failed")

        return bytes(into)

    def encrypt_inplace(
        self,
        key: bytes,
        nonce: bytes,
        buffer,  # bytearray or memoryview
        associated_data: bytes | None = None,
    ) -> bytes:
        """
        Encrypt buffer in-place, return authentication tag.

        This method encrypts data directly in the provided mutable buffer,
        avoiding memory allocation and reducing bandwidth usage. This is
        especially beneficial for large buffers (>10MB).

        Args:
            key: Encryption key
            nonce: Nonce
            buffer: Mutable buffer (bytearray or memoryview) containing plaintext.
                   Will be overwritten with ciphertext in-place.
            associated_data: Optional additional authenticated data

        Returns:
            Authentication tag (16 or 32 bytes depending on tag_size)

        Raises:
            TypeError: If buffer is not mutable (bytearray/memoryview)
            ValueError: If key or nonce has incorrect length
            AegisError: If encryption fails

        Example:
            >>> cipher = Aegis128X4()
            >>> key, nonce = cipher.random_key(), cipher.random_nonce()
            >>> buffer = bytearray(b"secret message")
            >>> tag = cipher.encrypt_inplace(key, nonce, buffer)
            >>> # buffer now contains ciphertext
            >>> # Send both: bytes(buffer) + tag

        Note:
            The plaintext is destroyed and replaced with ciphertext.
            For large buffers (>10MB), this can provide 30-50% performance
            improvement over regular encrypt() due to reduced memory bandwidth.
        """
        self._check_key(key)
        self._check_nonce(nonce)

        # Validate buffer type
        if not isinstance(buffer, (bytearray, memoryview)):
            raise TypeError(f"buffer must be bytearray or memoryview, got {type(buffer).__name__}")

        plaintext_len = len(buffer)

        # Handle associated data
        if associated_data is None:
            associated_data = b""

        assert self._encrypt_detached_func is not None

        # Allocate tag buffer (small, so overhead is negligible)
        tag_into = bytearray(self.tag_size)
        tag_buf = ffi.from_buffer(tag_into)

        # Get pointer to mutable buffer (for both input and output)
        buf_ptr = ffi.from_buffer(buffer)

        # Encrypt in-place: buf_ptr is both source and destination
        result = self._encrypt_detached_func(
            buf_ptr,  # Output ciphertext (overwrites input)
            tag_buf,
            self.tag_size,
            buf_ptr,  # Input plaintext (same as output)
            plaintext_len,
            associated_data if len(associated_data) > 0 else ffi.NULL,
            len(associated_data),
            nonce,
            key,
        )

        if result != 0:
            raise AegisError("Encryption failed")

        return bytes(tag_into)

    def decrypt_inplace(
        self,
        key: bytes,
        nonce: bytes,
        buffer,  # bytearray or memoryview
        tag: bytes,
        associated_data: bytes | None = None,
    ) -> None:
        """
        Decrypt buffer in-place and verify authentication tag.

        This method decrypts data directly in the provided mutable buffer,
        avoiding memory allocation and reducing bandwidth usage.

        Args:
            key: Encryption key
            nonce: Nonce
            buffer: Mutable buffer (bytearray or memoryview) containing ciphertext.
                   Will be overwritten with plaintext in-place.
            tag: Authentication tag
            associated_data: Optional additional authenticated data

        Raises:
            TypeError: If buffer is not mutable (bytearray/memoryview)
            ValueError: If key or nonce has incorrect length
            DecryptionError: If authentication fails

        Example:
            >>> cipher = Aegis128X4()
            >>> # Received ciphertext and tag
            >>> buffer = bytearray(ciphertext)
            >>> cipher.decrypt_inplace(key, nonce, buffer, tag)
            >>> # buffer now contains plaintext
            >>> plaintext = bytes(buffer)

        Note:
            If authentication fails, the buffer contents are zeroed for security.
            For large buffers (>10MB), this provides 30-50% performance improvement
            over regular decrypt() due to reduced memory bandwidth.
        """
        self._check_key(key)
        self._check_nonce(nonce)

        # Validate buffer type
        if not isinstance(buffer, (bytearray, memoryview)):
            raise TypeError(f"buffer must be bytearray or memoryview, got {type(buffer).__name__}")

        ciphertext_len = len(buffer)

        # Handle associated data
        if associated_data is None:
            associated_data = b""

        assert self._decrypt_detached_func is not None

        # Get pointer to mutable buffer (for both input and output)
        buf_ptr = ffi.from_buffer(buffer)

        # Decrypt in-place: buf_ptr is both source and destination
        result = self._decrypt_detached_func(
            buf_ptr,  # Output plaintext (overwrites input)
            buf_ptr,  # Input ciphertext (same as output)
            ciphertext_len,
            tag,
            len(tag),
            associated_data if len(associated_data) > 0 else ffi.NULL,
            len(associated_data),
            nonce,
            key,
        )

        if result != 0:
            # Zero buffer on authentication failure for security
            for i in range(len(buffer)):
                buffer[i] = 0
            raise DecryptionError("Authentication failed")


class _AEGISWithStream(_AEGISWithDetached):
    """Base class for AEGIS variants that support stream generation."""

    @classmethod
    def stream(cls, key: bytes, nonce: bytes, length: int, into: bytearray | None = None) -> bytes:
        """
        Generate a deterministic pseudo-random byte sequence.

        Args:
            key: Encryption key
            nonce: Nonce (can be reused for stream generation)
            length: Number of bytes to generate
            into: Optional pre-allocated bytearray for output (must be length bytes)

        Returns:
            Pseudo-random bytes

        Raises:
            ValueError: If key has incorrect length or into has wrong size
        """
        if len(key) != cls.KEY_SIZE:
            raise ValueError(f"Key must be {cls.KEY_SIZE} bytes, got {len(key)}")

        # Use pre-allocated buffer if provided
        if into is not None:
            if len(into) != length:
                raise ValueError(f"Output buffer must be {length} bytes, got {len(into)}")
            output_buf = ffi.from_buffer(into)
        else:
            into = bytearray(length)
            output_buf = ffi.from_buffer(into)

        assert cls._stream_func is not None
        cls._stream_func(output_buf, length, nonce, key)
        return bytes(into)


class Aegis128L(_AEGISWithStream):
    """
    AEGIS-128L authenticated encryption.

    Uses a 16-byte (128-bit) key and 16-byte nonce.
    Provides high performance on modern CPUs with AES-NI support.

    Example:
        >>> cipher = Aegis128L()
        >>> key = cipher.random_key()
        >>> nonce = cipher.random_nonce()
        >>> ciphertext = cipher.encrypt(key, nonce, b"secret message")
        >>> plaintext = cipher.decrypt(key, nonce, ciphertext)
    """

    KEY_SIZE = 16
    NONCE_SIZE = 16

    _encrypt_func = lib.aegis128l_encrypt
    _decrypt_func = lib.aegis128l_decrypt
    _encrypt_detached_func = lib.aegis128l_encrypt_detached
    _decrypt_detached_func = lib.aegis128l_decrypt_detached
    _stream_func = lib.aegis128l_stream


class Aegis256(_AEGISWithStream):
    """
    AEGIS-256 authenticated encryption.

    Uses a 32-byte (256-bit) key and 32-byte nonce.
    Provides higher security margin than AEGIS-128L.

    Example:
        >>> cipher = Aegis256()
        >>> key = cipher.random_key()
        >>> nonce = cipher.random_nonce()
        >>> ciphertext = cipher.encrypt(key, nonce, b"secret message")
        >>> plaintext = cipher.decrypt(key, nonce, ciphertext)
    """

    KEY_SIZE = 32
    NONCE_SIZE = 32

    _encrypt_func = lib.aegis256_encrypt
    _decrypt_func = lib.aegis256_decrypt
    _encrypt_detached_func = lib.aegis256_encrypt_detached
    _decrypt_detached_func = lib.aegis256_decrypt_detached
    _stream_func = lib.aegis256_stream


class Aegis128X2(_AEGISWithDetached):
    """
    AEGIS-128X2 - dual-lane variant for higher performance.

    Uses a 16-byte key and 16-byte nonce.
    Provides higher throughput on CPUs with wide SIMD capabilities.

    Example:
        >>> cipher = Aegis128X2()
        >>> key = cipher.random_key()
        >>> nonce = cipher.random_nonce()
        >>> ciphertext = cipher.encrypt(key, nonce, b"secret message")
        >>> plaintext = cipher.decrypt(key, nonce, ciphertext)
    """

    KEY_SIZE = 16
    NONCE_SIZE = 16

    _encrypt_detached_func = lib.aegis128x2_encrypt_detached
    _decrypt_detached_func = lib.aegis128x2_decrypt_detached


class Aegis128X4(_AEGISWithDetached):
    """
    AEGIS-128X4 - quad-lane variant for highest performance on AVX-512.

    Uses a 16-byte key and 16-byte nonce.
    Provides maximum throughput on CPUs with AVX-512 support.

    Example:
        >>> cipher = Aegis128X4()
        >>> key = cipher.random_key()
        >>> nonce = cipher.random_nonce()
        >>> ciphertext = cipher.encrypt(key, nonce, b"secret message")
        >>> plaintext = cipher.decrypt(key, nonce, ciphertext)
    """

    KEY_SIZE = 16
    NONCE_SIZE = 16

    _encrypt_detached_func = lib.aegis128x4_encrypt_detached
    _decrypt_detached_func = lib.aegis128x4_decrypt_detached


class Aegis256X2(_AEGISWithDetached):
    """
    AEGIS-256X2 - dual-lane variant with 256-bit security.

    Uses a 32-byte key and 32-byte nonce.
    Provides higher throughput with increased security margin.

    Example:
        >>> cipher = Aegis256X2()
        >>> key = cipher.random_key()
        >>> nonce = cipher.random_nonce()
        >>> ciphertext = cipher.encrypt(key, nonce, b"secret message")
        >>> plaintext = cipher.decrypt(key, nonce, ciphertext)
    """

    KEY_SIZE = 32
    NONCE_SIZE = 32

    _encrypt_detached_func = lib.aegis256x2_encrypt_detached
    _decrypt_detached_func = lib.aegis256x2_decrypt_detached


class Aegis256X4(_AEGISWithDetached):
    """
    AEGIS-256X4 - quad-lane variant with 256-bit security.

    Uses a 32-byte key and 32-byte nonce.
    Provides maximum throughput with increased security margin on AVX-512 CPUs.

    Example:
        >>> cipher = Aegis256X4()
        >>> key = cipher.random_key()
        >>> nonce = cipher.random_nonce()
        >>> ciphertext = cipher.encrypt(key, nonce, b"secret message")
        >>> plaintext = cipher.decrypt(key, nonce, ciphertext)
    """

    KEY_SIZE = 32
    NONCE_SIZE = 32

    _encrypt_detached_func = lib.aegis256x4_encrypt_detached
    _decrypt_detached_func = lib.aegis256x4_decrypt_detached


# MAC classes for message authentication


class _AEGISMACBase:
    """Base class for AEGIS-MAC implementations with shared functionality."""

    KEY_SIZE: int
    NONCE_SIZE: int
    TAG_SIZE_MIN: int = 16
    TAG_SIZE_MAX: int = 32

    # Subclasses should define these
    _mac_state_type = None
    _mac_init_func = None
    _mac_update_func = None
    _mac_final_func = None
    _mac_verify_func = None
    _mac_reset_func = None

    def __init__(self, key: bytes, nonce: bytes, tag_size: int = 32) -> None:
        """
        Initialize an AEGIS-MAC instance.

        Args:
            key: MAC key (size depends on cipher variant)
            nonce: Nonce (size depends on cipher variant)
            tag_size: Size of the authentication tag in bytes (16 or 32).
                     32 bytes is recommended for maximum security.

        Raises:
            ValueError: If key, nonce, or tag_size is invalid

        Note:
            The same key MUST NOT be used both for MAC and encryption.
            Generate a random key and keep it secret for secure MAC operations.
        """
        if len(key) != self.KEY_SIZE:
            raise ValueError(f"Key must be {self.KEY_SIZE} bytes, got {len(key)}")
        if len(nonce) != self.NONCE_SIZE:
            raise ValueError(f"Nonce must be {self.NONCE_SIZE} bytes, got {len(nonce)}")
        if tag_size not in (16, 32):
            raise ValueError(f"tag_size must be 16 or 32, got {tag_size}")

        self.tag_size = tag_size
        self._state = ffi.new(f"{self._mac_state_type}*")
        self._mac_init_func(self._state, key, nonce)
        self._finalized = False

    @classmethod
    def random_key(cls) -> bytes:
        """Generate a random key suitable for this MAC."""
        return os.urandom(cls.KEY_SIZE)

    @classmethod
    def random_nonce(cls) -> bytes:
        """Generate a random nonce suitable for this MAC."""
        return os.urandom(cls.NONCE_SIZE)

    def update(self, data: bytes) -> None:
        """
        Update the MAC state with input data.

        Args:
            data: Input data to authenticate

        Raises:
            AegisError: If update fails or MAC has been finalized
        """
        if self._finalized:
            raise AegisError("Cannot update after finalization")

        result = self._mac_update_func(self._state, data, len(data))
        if result != 0:
            raise AegisError("MAC update failed")

    def final(self) -> bytes:
        """
        Finalize the MAC and generate the authentication tag.

        Returns:
            Authentication tag

        Raises:
            AegisError: If MAC has already been finalized
        """
        if self._finalized:
            raise AegisError("MAC already finalized")

        tag_buf = ffi.new(f"uint8_t[{self.tag_size}]")
        result = self._mac_final_func(self._state, tag_buf, self.tag_size)
        if result != 0:
            raise AegisError("MAC finalization failed")

        self._finalized = True
        return bytes(ffi.buffer(tag_buf, self.tag_size))

    def verify(self, tag: bytes) -> None:
        """
        Verify a MAC in constant time.

        Args:
            tag: Authentication tag to verify

        Raises:
            DecryptionError: If the tag is not authentic
            AegisError: If MAC has already been finalized
        """
        if self._finalized:
            raise AegisError("MAC already finalized")

        result = self._mac_verify_func(self._state, tag, len(tag))
        self._finalized = True

        if result != 0:
            raise DecryptionError("MAC verification failed")

    def reset(self) -> None:
        """Reset the MAC state for reuse."""
        self._mac_reset_func(self._state)
        self._finalized = False


class AegisMac128L(_AEGISMACBase):
    """
    AEGIS-128L MAC - message authentication using AEGIS-128L.

    Uses a 16-byte key and 16-byte nonce.

    Example:
        >>> mac = AegisMac128L(key, nonce)
        >>> mac.update(b"hello")
        >>> mac.update(b" world")
        >>> tag = mac.final()

    Example (verification):
        >>> mac2 = AegisMac128L(key, nonce)
        >>> mac2.update(b"hello world")
        >>> mac2.verify(tag)  # raises DecryptionError if invalid
    """

    KEY_SIZE = 16
    NONCE_SIZE = 16

    _mac_state_type = "aegis128l_mac_state"
    _mac_init_func = lib.aegis128l_mac_init
    _mac_update_func = lib.aegis128l_mac_update
    _mac_final_func = lib.aegis128l_mac_final
    _mac_verify_func = lib.aegis128l_mac_verify
    _mac_reset_func = lib.aegis128l_mac_reset


class AegisMac256(_AEGISMACBase):
    """
    AEGIS-256 MAC - message authentication using AEGIS-256.

    Uses a 32-byte key and 32-byte nonce.

    Example:
        >>> mac = AegisMac256(key, nonce)
        >>> mac.update(b"hello")
        >>> mac.update(b" world")
        >>> tag = mac.final()
    """

    KEY_SIZE = 32
    NONCE_SIZE = 32

    _mac_state_type = "aegis256_mac_state"
    _mac_init_func = lib.aegis256_mac_init
    _mac_update_func = lib.aegis256_mac_update
    _mac_final_func = lib.aegis256_mac_final
    _mac_verify_func = lib.aegis256_mac_verify
    _mac_reset_func = lib.aegis256_mac_reset


class AegisMac128X2(_AEGISMACBase):
    """
    AEGIS-128X2 MAC - message authentication using dual-lane AEGIS-128.

    Uses a 16-byte key and 16-byte nonce.
    Provides higher throughput on CPUs with wide SIMD capabilities.

    Example:
        >>> mac = AegisMac128X2(key, nonce)
        >>> mac.update(b"data")
        >>> tag = mac.final()
    """

    KEY_SIZE = 16
    NONCE_SIZE = 16

    _mac_state_type = "aegis128x2_mac_state"
    _mac_init_func = lib.aegis128x2_mac_init
    _mac_update_func = lib.aegis128x2_mac_update
    _mac_final_func = lib.aegis128x2_mac_final
    _mac_verify_func = lib.aegis128x2_mac_verify
    _mac_reset_func = lib.aegis128x2_mac_reset


class AegisMac128X4(_AEGISMACBase):
    """
    AEGIS-128X4 MAC - message authentication using quad-lane AEGIS-128.

    Uses a 16-byte key and 16-byte nonce.
    Provides maximum throughput on CPUs with AVX-512 support.

    Example:
        >>> mac = AegisMac128X4(key, nonce)
        >>> mac.update(b"data")
        >>> tag = mac.final()
    """

    KEY_SIZE = 16
    NONCE_SIZE = 16

    _mac_state_type = "aegis128x4_mac_state"
    _mac_init_func = lib.aegis128x4_mac_init
    _mac_update_func = lib.aegis128x4_mac_update
    _mac_final_func = lib.aegis128x4_mac_final
    _mac_verify_func = lib.aegis128x4_mac_verify
    _mac_reset_func = lib.aegis128x4_mac_reset


class AegisMac256X2(_AEGISMACBase):
    """
    AEGIS-256X2 MAC - message authentication using dual-lane AEGIS-256.

    Uses a 32-byte key and 32-byte nonce.
    Provides higher throughput with increased security margin.

    Example:
        >>> mac = AegisMac256X2(key, nonce)
        >>> mac.update(b"data")
        >>> tag = mac.final()
    """

    KEY_SIZE = 32
    NONCE_SIZE = 32

    _mac_state_type = "aegis256x2_mac_state"
    _mac_init_func = lib.aegis256x2_mac_init
    _mac_update_func = lib.aegis256x2_mac_update
    _mac_final_func = lib.aegis256x2_mac_final
    _mac_verify_func = lib.aegis256x2_mac_verify
    _mac_reset_func = lib.aegis256x2_mac_reset


class AegisMac256X4(_AEGISMACBase):
    """
    AEGIS-256X4 MAC - message authentication using quad-lane AEGIS-256.

    Uses a 32-byte key and 32-byte nonce.
    Provides maximum throughput with increased security margin on AVX-512 CPUs.

    Example:
        >>> mac = AegisMac256X4(key, nonce)
        >>> mac.update(b"data")
        >>> tag = mac.final()
    """

    KEY_SIZE = 32
    NONCE_SIZE = 32

    _mac_state_type = "aegis256x4_mac_state"
    _mac_init_func = lib.aegis256x4_mac_init
    _mac_update_func = lib.aegis256x4_mac_update
    _mac_final_func = lib.aegis256x4_mac_final
    _mac_verify_func = lib.aegis256x4_mac_verify
    _mac_reset_func = lib.aegis256x4_mac_reset


class _AEGISStreamEncryptBase:
    """Base class for AEGIS streaming encryption with shared functionality."""

    KEY_SIZE: int
    NONCE_SIZE: int
    TAG_SIZE_MIN: int = 16
    TAG_SIZE_MAX: int = 32

    # Subclasses should define these
    _state_type = None
    _state_init_func = None
    _state_encrypt_update_func = None
    _state_encrypt_final_func = None

    def __init__(
        self,
        key: bytes,
        nonce: bytes,
        associated_data: bytes | None = None,
        tag_size: int = 32,
    ) -> None:
        """
        Initialize a streaming encryption context.

        Args:
            key: Encryption key (size depends on cipher variant)
            nonce: Nonce (must be unique for each message with the same key)
            associated_data: Optional additional authenticated data (not encrypted)
            tag_size: Size of the authentication tag in bytes (16 or 32).
                     32 bytes is recommended for maximum security.

        Raises:
            ValueError: If key, nonce, or tag_size is invalid

        Note:
            After initialization, call update() with plaintext chunks, then final()
            to get the authentication tag. The ciphertext is produced incrementally.
        """
        if len(key) != self.KEY_SIZE:
            raise ValueError(f"Key must be {self.KEY_SIZE} bytes, got {len(key)}")
        if len(nonce) != self.NONCE_SIZE:
            raise ValueError(f"Nonce must be {self.NONCE_SIZE} bytes, got {len(nonce)}")
        if tag_size not in (16, 32):
            raise ValueError(f"tag_size must be 16 or 32, got {tag_size}")

        self.tag_size = tag_size
        self._state = ffi.new(f"{self._state_type}*")

        ad = associated_data or b""
        self._state_init_func(self._state, ad, len(ad), nonce, key)
        self._finalized = False

    @classmethod
    def random_key(cls) -> bytes:
        """Generate a random key suitable for this cipher."""
        return os.urandom(cls.KEY_SIZE)

    @classmethod
    def random_nonce(cls) -> bytes:
        """Generate a random nonce suitable for this cipher."""
        return os.urandom(cls.NONCE_SIZE)

    def update(self, plaintext: bytes) -> bytes:
        """
        Encrypt a chunk of plaintext.

        Args:
            plaintext: Plaintext data to encrypt

        Returns:
            Ciphertext chunk (same length as plaintext)

        Raises:
            AegisError: If update fails or encryption has been finalized
        """
        if self._finalized:
            raise AegisError("Cannot update after finalization")

        if not plaintext:
            return b""

        # Allocate output buffer (same size as input - immediate output)
        ciphertext_buf = ffi.new(f"uint8_t[{len(plaintext)}]")

        result = self._state_encrypt_update_func(
            self._state, ciphertext_buf, plaintext, len(plaintext)
        )

        if result != 0:
            raise AegisError("Encryption update failed")

        return bytes(ffi.buffer(ciphertext_buf, len(plaintext)))

    def final(self) -> bytes:
        """
        Finalize the encryption and generate the authentication tag.

        Returns:
            Authentication tag

        Raises:
            AegisError: If encryption has already been finalized
        """
        if self._finalized:
            raise AegisError("Encryption already finalized")

        tag_buf = ffi.new(f"uint8_t[{self.tag_size}]")

        result = self._state_encrypt_final_func(self._state, tag_buf, self.tag_size)

        if result != 0:
            raise AegisError("Encryption finalization failed")

        self._finalized = True
        return bytes(ffi.buffer(tag_buf, self.tag_size))

    def __enter__(self):
        """Context manager entry."""
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit - finalize if not already done."""
        if not self._finalized:
            self.final()
        return False


class _AEGISStreamDecryptBase:
    """Base class for AEGIS streaming decryption with shared functionality."""

    KEY_SIZE: int
    NONCE_SIZE: int
    TAG_SIZE_MIN: int = 16
    TAG_SIZE_MAX: int = 32

    # Subclasses should define these
    _state_type = None
    _state_init_func = None
    _state_decrypt_update_func = None
    _state_decrypt_final_func = None

    def __init__(
        self,
        key: bytes,
        nonce: bytes,
        associated_data: bytes | None = None,
        tag_size: int = 32,
    ) -> None:
        """
        Initialize a streaming decryption context.

        Args:
            key: Decryption key (size depends on cipher variant)
            nonce: Nonce (must match the one used for encryption)
            associated_data: Optional additional authenticated data (must match encryption)
            tag_size: Size of the authentication tag in bytes (16 or 32).

        Raises:
            ValueError: If key, nonce, or tag_size is invalid

        Note:
            After initialization, call update() with ciphertext chunks, then verify()
            with the authentication tag. The plaintext is produced incrementally but
            MUST NOT be released until the tag has been verified.
        """
        if len(key) != self.KEY_SIZE:
            raise ValueError(f"Key must be {self.KEY_SIZE} bytes, got {len(key)}")
        if len(nonce) != self.NONCE_SIZE:
            raise ValueError(f"Nonce must be {self.NONCE_SIZE} bytes, got {len(nonce)}")
        if tag_size not in (16, 32):
            raise ValueError(f"tag_size must be 16 or 32, got {tag_size}")

        self.tag_size = tag_size
        self._state = ffi.new(f"{self._state_type}*")

        ad = associated_data or b""
        self._state_init_func(self._state, ad, len(ad), nonce, key)
        self._finalized = False
        self._plaintext_chunks = []

    @classmethod
    def random_key(cls) -> bytes:
        """Generate a random key suitable for this cipher."""
        return os.urandom(cls.KEY_SIZE)

    @classmethod
    def random_nonce(cls) -> bytes:
        """Generate a random nonce suitable for this cipher."""
        return os.urandom(cls.NONCE_SIZE)

    def update(self, ciphertext: bytes) -> None:
        """
        Decrypt a chunk of ciphertext.

        SECURITY WARNING: Do not release the plaintext until verify() succeeds.
        This implementation buffers the plaintext internally and only releases
        it after successful verification.

        Args:
            ciphertext: Ciphertext data to decrypt

        Raises:
            AegisError: If update fails or decryption has been finalized
        """
        if self._finalized:
            raise AegisError("Cannot update after finalization")

        if not ciphertext:
            return

        # Allocate output buffer (same size as input - immediate output)
        plaintext_buf = ffi.new(f"uint8_t[{len(ciphertext)}]")

        result = self._state_decrypt_update_func(
            self._state, plaintext_buf, ciphertext, len(ciphertext)
        )

        if result != 0:
            raise AegisError("Decryption update failed")

        # Buffer the plaintext - don't release it until verification
        self._plaintext_chunks.append(bytes(ffi.buffer(plaintext_buf, len(ciphertext))))

    def verify(self, tag: bytes) -> bytes:
        """
        Verify the authentication tag and return all decrypted plaintext.

        Args:
            tag: Authentication tag to verify

        Returns:
            Complete plaintext (all chunks concatenated)

        Raises:
            DecryptionError: If the tag is not authentic
            AegisError: If decryption has already been finalized
        """
        if self._finalized:
            raise AegisError("Decryption already finalized")

        result = self._state_decrypt_final_func(self._state, tag, len(tag))

        self._finalized = True

        if result != 0:
            # Clear plaintext on verification failure for security
            self._plaintext_chunks.clear()
            raise DecryptionError("Tag verification failed")

        # Verification succeeded - return the plaintext
        plaintext = b"".join(self._plaintext_chunks)
        self._plaintext_chunks.clear()
        return plaintext

    def __enter__(self):
        """Context manager entry."""
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit - clear buffered plaintext."""
        self._plaintext_chunks.clear()
        return False


# AEGIS-128L Streaming Classes


class AegisStreamEncrypt128L(_AEGISStreamEncryptBase):
    """
    AEGIS-128L streaming encryption.

    Uses a 16-byte key and 16-byte nonce.
    Provides high throughput with 256-bit security.

    Example:
        >>> key = AegisStreamEncrypt128L.random_key()
        >>> nonce = AegisStreamEncrypt128L.random_nonce()
        >>> enc = AegisStreamEncrypt128L(key, nonce, b"metadata")
        >>> c1 = enc.update(b"hello ")
        >>> c2 = enc.update(b"world")
        >>> tag = enc.final()

    Example (with context manager):
        >>> with AegisStreamEncrypt128L(key, nonce) as enc:
        ...     ciphertext = enc.update(plaintext)
        ...     tag = enc.final()
    """

    KEY_SIZE = 16
    NONCE_SIZE = 16

    _state_type = "aegis128l_state"
    _state_init_func = lib.aegis128l_state_init
    _state_encrypt_update_func = lib.aegis128l_state_encrypt_update
    _state_encrypt_final_func = lib.aegis128l_state_encrypt_final


class AegisStreamDecrypt128L(_AEGISStreamDecryptBase):
    """
    AEGIS-128L streaming decryption.

    Uses a 16-byte key and 16-byte nonce.

    Example:
        >>> dec = AegisStreamDecrypt128L(key, nonce, b"metadata")
        >>> dec.update(ciphertext1)
        >>> dec.update(ciphertext2)
        >>> plaintext = dec.verify(tag)  # raises DecryptionError if invalid

    Example (with context manager):
        >>> with AegisStreamDecrypt128L(key, nonce) as dec:
        ...     dec.update(ciphertext)
        ...     plaintext = dec.verify(tag)
    """

    KEY_SIZE = 16
    NONCE_SIZE = 16

    _state_type = "aegis128l_state"
    _state_init_func = lib.aegis128l_state_init
    _state_decrypt_update_func = lib.aegis128l_state_decrypt_update
    _state_decrypt_final_func = lib.aegis128l_state_decrypt_final


# AEGIS-256 Streaming Classes


class AegisStreamEncrypt256(_AEGISStreamEncryptBase):
    """
    AEGIS-256 streaming encryption.

    Uses a 32-byte key and 32-byte nonce.
    Provides maximum security with 256-bit key.

    Example:
        >>> key = AegisStreamEncrypt256.random_key()
        >>> nonce = AegisStreamEncrypt256.random_nonce()
        >>> enc = AegisStreamEncrypt256(key, nonce)
        >>> ciphertext = enc.update(plaintext)
        >>> tag = enc.final()
    """

    KEY_SIZE = 32
    NONCE_SIZE = 32

    _state_type = "aegis256_state"
    _state_init_func = lib.aegis256_state_init
    _state_encrypt_update_func = lib.aegis256_state_encrypt_update
    _state_encrypt_final_func = lib.aegis256_state_encrypt_final


class AegisStreamDecrypt256(_AEGISStreamDecryptBase):
    """
    AEGIS-256 streaming decryption.

    Uses a 32-byte key and 32-byte nonce.

    Example:
        >>> dec = AegisStreamDecrypt256(key, nonce)
        >>> dec.update(ciphertext)
        >>> plaintext = dec.verify(tag)
    """

    KEY_SIZE = 32
    NONCE_SIZE = 32

    _state_type = "aegis256_state"
    _state_init_func = lib.aegis256_state_init
    _state_decrypt_update_func = lib.aegis256_state_decrypt_update
    _state_decrypt_final_func = lib.aegis256_state_decrypt_final


# AEGIS-128X2 Streaming Classes


class AegisStreamEncrypt128X2(_AEGISStreamEncryptBase):
    """
    AEGIS-128X2 streaming encryption - dual-lane variant.

    Uses a 16-byte key and 16-byte nonce.
    Provides higher throughput on CPUs with SIMD support.

    Example:
        >>> enc = AegisStreamEncrypt128X2(key, nonce)
        >>> ciphertext = enc.update(plaintext)
        >>> tag = enc.final()
    """

    KEY_SIZE = 16
    NONCE_SIZE = 16

    _state_type = "aegis128x2_state"
    _state_init_func = lib.aegis128x2_state_init
    _state_encrypt_update_func = lib.aegis128x2_state_encrypt_update
    _state_encrypt_final_func = lib.aegis128x2_state_encrypt_final


class AegisStreamDecrypt128X2(_AEGISStreamDecryptBase):
    """
    AEGIS-128X2 streaming decryption - dual-lane variant.

    Uses a 16-byte key and 16-byte nonce.

    Example:
        >>> dec = AegisStreamDecrypt128X2(key, nonce)
        >>> dec.update(ciphertext)
        >>> plaintext = dec.verify(tag)
    """

    KEY_SIZE = 16
    NONCE_SIZE = 16

    _state_type = "aegis128x2_state"
    _state_init_func = lib.aegis128x2_state_init
    _state_decrypt_update_func = lib.aegis128x2_state_decrypt_update
    _state_decrypt_final_func = lib.aegis128x2_state_decrypt_final


# AEGIS-128X4 Streaming Classes


class AegisStreamEncrypt128X4(_AEGISStreamEncryptBase):
    """
    AEGIS-128X4 streaming encryption - quad-lane variant.

    Uses a 16-byte key and 16-byte nonce.
    Provides maximum throughput on CPUs with wide SIMD support.

    Example:
        >>> enc = AegisStreamEncrypt128X4(key, nonce)
        >>> ciphertext = enc.update(plaintext)
        >>> tag = enc.final()
    """

    KEY_SIZE = 16
    NONCE_SIZE = 16

    _state_type = "aegis128x4_state"
    _state_init_func = lib.aegis128x4_state_init
    _state_encrypt_update_func = lib.aegis128x4_state_encrypt_update
    _state_encrypt_final_func = lib.aegis128x4_state_encrypt_final


class AegisStreamDecrypt128X4(_AEGISStreamDecryptBase):
    """
    AEGIS-128X4 streaming decryption - quad-lane variant.

    Uses a 16-byte key and 16-byte nonce.

    Example:
        >>> dec = AegisStreamDecrypt128X4(key, nonce)
        >>> dec.update(ciphertext)
        >>> plaintext = dec.verify(tag)
    """

    KEY_SIZE = 16
    NONCE_SIZE = 16

    _state_type = "aegis128x4_state"
    _state_init_func = lib.aegis128x4_state_init
    _state_decrypt_update_func = lib.aegis128x4_state_decrypt_update
    _state_decrypt_final_func = lib.aegis128x4_state_decrypt_final


# AEGIS-256X2 Streaming Classes


class AegisStreamEncrypt256X2(_AEGISStreamEncryptBase):
    """
    AEGIS-256X2 streaming encryption - dual-lane 256-bit variant.

    Uses a 32-byte key and 32-byte nonce.
    Provides high throughput with maximum security.

    Example:
        >>> enc = AegisStreamEncrypt256X2(key, nonce)
        >>> ciphertext = enc.update(plaintext)
        >>> tag = enc.final()
    """

    KEY_SIZE = 32
    NONCE_SIZE = 32

    _state_type = "aegis256x2_state"
    _state_init_func = lib.aegis256x2_state_init
    _state_encrypt_update_func = lib.aegis256x2_state_encrypt_update
    _state_encrypt_final_func = lib.aegis256x2_state_encrypt_final


class AegisStreamDecrypt256X2(_AEGISStreamDecryptBase):
    """
    AEGIS-256X2 streaming decryption - dual-lane 256-bit variant.

    Uses a 32-byte key and 32-byte nonce.

    Example:
        >>> dec = AegisStreamDecrypt256X2(key, nonce)
        >>> dec.update(ciphertext)
        >>> plaintext = dec.verify(tag)
    """

    KEY_SIZE = 32
    NONCE_SIZE = 32

    _state_type = "aegis256x2_state"
    _state_init_func = lib.aegis256x2_state_init
    _state_decrypt_update_func = lib.aegis256x2_state_decrypt_update
    _state_decrypt_final_func = lib.aegis256x2_state_decrypt_final


# AEGIS-256X4 Streaming Classes


class AegisStreamEncrypt256X4(_AEGISStreamEncryptBase):
    """
    AEGIS-256X4 streaming encryption - quad-lane 256-bit variant.

    Uses a 32-byte key and 32-byte nonce.
    Provides maximum throughput and security on AVX-512 CPUs.

    Example:
        >>> enc = AegisStreamEncrypt256X4(key, nonce)
        >>> ciphertext = enc.update(plaintext)
        >>> tag = enc.final()
    """

    KEY_SIZE = 32
    NONCE_SIZE = 32

    _state_type = "aegis256x4_state"
    _state_init_func = lib.aegis256x4_state_init
    _state_encrypt_update_func = lib.aegis256x4_state_encrypt_update
    _state_encrypt_final_func = lib.aegis256x4_state_encrypt_final


class AegisStreamDecrypt256X4(_AEGISStreamDecryptBase):
    """
    AEGIS-256X4 streaming decryption - quad-lane 256-bit variant.

    Uses a 32-byte key and 32-byte nonce.

    Example:
        >>> dec = AegisStreamDecrypt256X4(key, nonce)
        >>> dec.update(ciphertext)
        >>> plaintext = dec.verify(tag)
    """

    KEY_SIZE = 32
    NONCE_SIZE = 32

    _state_type = "aegis256x4_state"
    _state_init_func = lib.aegis256x4_state_init
    _state_decrypt_update_func = lib.aegis256x4_state_decrypt_update
    _state_decrypt_final_func = lib.aegis256x4_state_decrypt_final
