"""
AES-256-GCM encryption for agent memories.

This module provides authenticated encryption for sensitive agent data using
AES-256 in Galois/Counter Mode (GCM). Each encryption uses a unique IV to
ensure semantic security.

Requirements: SEC-01 (encryption for agent memories)
"""

from __future__ import annotations

import os
from dataclasses import dataclass, field
from typing import Optional

from cryptography.hazmat.primitives.ciphers.aead import AESGCM
import cryptography.exceptions


@dataclass
class EncryptedData:
    """
    Container for encrypted data with integrity verification.

    AES-GCM combines encryption and authentication, producing ciphertext
    that includes an authentication tag. The IV (Initialization Vector)
    must be unique for each encryption with the same key.

    Attributes:
        iv: 12-byte initialization vector (96 bits, recommended for GCM).
        ciphertext: Encrypted data including the 16-byte authentication tag.
        associated_data: Optional context for authenticated but unencrypted data.
    """

    iv: bytes
    ciphertext: bytes
    associated_data: Optional[bytes] = None

    def __post_init__(self):
        """Validate encrypted data structure."""
        if len(self.iv) != MemoryEncryptor.IV_SIZE:
            raise ValueError(
                f"IV must be {MemoryEncryptor.IV_SIZE} bytes, got {len(self.iv)}"
            )

    def to_dict(self) -> dict:
        """
        Convert to dictionary for serialization.

        Returns:
            Dictionary with base64-encoded binary fields.
        """
        import base64

        return {
            "iv": base64.b64encode(self.iv).decode("ascii"),
            "ciphertext": base64.b64encode(self.ciphertext).decode("ascii"),
            "associated_data": (
                base64.b64encode(self.associated_data).decode("ascii")
                if self.associated_data
                else None
            ),
        }

    @classmethod
    def from_dict(cls, data: dict) -> "EncryptedData":
        """
        Create EncryptedData from dictionary.

        Args:
            data: Dictionary with base64-encoded binary fields.

        Returns:
            EncryptedData instance.
        """
        import base64

        return cls(
            iv=base64.b64decode(data["iv"]),
            ciphertext=base64.b64decode(data["ciphertext"]),
            associated_data=(
                base64.b64decode(data["associated_data"])
                if data.get("associated_data")
                else None
            ),
        )


class MemoryEncryptor:
    """
    AES-256-GCM encryptor for agent memories.

    Provides authenticated encryption with the following properties:
    - Confidentiality: Data cannot be read without the key
    - Integrity: Any tampering is detected on decryption
    - Authenticity: Encryption key is required to decrypt

    Security notes:
    - Never reuse an IV with the same key
    - Keep the encryption key secure
    - Use associated_data for context binding

    Example:
        >>> encryptor = MemoryEncryptor()
        >>> encrypted = encryptor.encrypt(b"secret memory")
        >>> decrypted = encryptor.decrypt(encrypted)
        >>> assert decrypted == b"secret memory"
    """

    KEY_SIZE = 32  # 256 bits for AES-256
    IV_SIZE = 12  # 96 bits (recommended for GCM)

    def __init__(self, key: Optional[bytes] = None):
        """
        Initialize encryptor with existing or generated key.

        Args:
            key: 32-byte encryption key. If None, generates a new key.
        """
        if key is None:
            key = self.generate_key()
        elif len(key) != self.KEY_SIZE:
            raise ValueError(f"Key must be {self.KEY_SIZE} bytes, got {len(key)}")

        self._key = key
        self._aesgcm = AESGCM(key)

    @classmethod
    def generate_key(cls) -> bytes:
        """
        Generate a new random encryption key.

        Uses cryptographically secure random number generation.

        Returns:
            32-byte (256-bit) encryption key.
        """
        return os.urandom(cls.KEY_SIZE)

    @property
    def key(self) -> bytes:
        """
        Get the encryption key.

        WARNING: Store this key securely. Losing it means losing all
        encrypted data. Sharing it compromises all encrypted data.

        Returns:
            The 32-byte encryption key.
        """
        return self._key

    def encrypt(
        self, plaintext: bytes, associated_data: Optional[bytes] = None
    ) -> EncryptedData:
        """
        Encrypt plaintext with AES-256-GCM.

        Generates a unique random IV for each encryption to ensure
        semantic security. The same plaintext encrypted twice will
        produce different ciphertexts.

        Args:
            plaintext: Data to encrypt.
            associated_data: Optional context data that is authenticated
                but not encrypted. Useful for binding ciphertext to
                a specific context (e.g., agent_id).

        Returns:
            EncryptedData containing IV, ciphertext, and associated data.

        Example:
            >>> encrypted = encryptor.encrypt(
            ...     b"secret",
            ...     associated_data=b"agent:researcher"
            ... )
        """
        # CRITICAL: Generate unique IV for each encryption
        iv = os.urandom(self.IV_SIZE)

        # AESGCM.encrypt returns ciphertext with authentication tag appended
        ciphertext = self._aesgcm.encrypt(iv, plaintext, associated_data)

        return EncryptedData(
            iv=iv,
            ciphertext=ciphertext,
            associated_data=associated_data,
        )

    def decrypt(self, encrypted: EncryptedData) -> bytes:
        """
        Decrypt ciphertext with AES-256-GCM.

        Verifies integrity and authenticity before returning plaintext.
        Raises InvalidTag if the data has been tampered with or if
        the wrong key is used.

        Args:
            encrypted: EncryptedData to decrypt.

        Returns:
            Decrypted plaintext bytes.

        Raises:
            cryptography.exceptions.InvalidTag: If ciphertext is corrupted,
                tampered, or decrypted with wrong key.

        Example:
            >>> encrypted = encryptor.encrypt(b"secret")
            >>> plaintext = encryptor.decrypt(encrypted)
            >>> assert plaintext == b"secret"
        """
        return self._aesgcm.decrypt(
            encrypted.iv,
            encrypted.ciphertext,
            encrypted.associated_data,
        )
