"""Encryption utilities for data at rest.

This module provides encryption/decryption capabilities for sensitive data
stored on disk, using Fernet (symmetric encryption) from the cryptography library.
"""

import base64
import hashlib
import logging
import os
import threading
import time
from typing import List, Optional

from cryptography.fernet import Fernet, InvalidToken, MultiFernet
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC

from ..utils.exceptions import ConfigurationError

logger = logging.getLogger(__name__)


class EncryptionManager:
    """Manages encryption and decryption of data at rest.

    Uses Fernet (symmetric encryption) with AES-128 in CBC mode.
    Supports key derivation from passwords using PBKDF2.

    Examples:
        >>> # Using a generated key
        >>> key = EncryptionManager.generate_key()
        >>> manager = EncryptionManager(key)
        >>> encrypted = manager.encrypt(b"sensitive data")
        >>> decrypted = manager.decrypt(encrypted)

        >>> # Using a password-derived key
        >>> manager = EncryptionManager.from_password("my-secret-password")
        >>> encrypted = manager.encrypt(b"sensitive data")
    """

    def __init__(self, key: bytes, salt: Optional[bytes] = None):
        """Initialize encryption manager with a key.

        Args:
            key: 32-byte Fernet key (base64-encoded)
            salt: Optional salt used for key derivation (for password-based keys)

        Raises:
            ValueError: If key is invalid
        """
        try:
            self._fernet = Fernet(key)
            self._key = key
            self._salt = salt  # Store salt for password-based keys
        except Exception as e:
            raise ValueError(f"Invalid encryption key: {e}")

    @property
    def salt(self) -> Optional[bytes]:
        """Get the salt used for password-based key derivation.

        Returns:
            Salt bytes if created via from_password(), None otherwise

        Note:
            You MUST save this salt to decrypt data later with the same password.
        """
        return self._salt

    @staticmethod
    def generate_key() -> bytes:
        """Generate a new Fernet encryption key.

        Returns:
            32-byte base64-encoded encryption key

        Example:
            >>> key = EncryptionManager.generate_key()
            >>> # Store this key securely (e.g., in environment variable or secret manager)
        """
        return Fernet.generate_key()

    @classmethod
    def from_password(
        cls, password: str, salt: Optional[bytes] = None, iterations: int = 480000
    ) -> "EncryptionManager":
        """Create encryption manager from a password.

        Derives a Fernet key from the password using PBKDF2-HMAC-SHA256.

        Args:
            password: Password to derive key from
            salt: Salt for key derivation (default: generates random salt)
            iterations: Number of PBKDF2 iterations (default: 480000, OWASP 2023 recommendation)

        Returns:
            EncryptionManager instance with salt stored on .salt property

        Important:
            The salt is automatically stored on the returned instance's .salt property.
            You MUST save this salt to persistent storage to be able to decrypt later.

        Example:
            >>> manager = EncryptionManager.from_password("my-password")
            >>> encrypted = manager.encrypt(b"data")
            >>> # CRITICAL: Save manager.salt for later use!
            >>> with open("salt.bin", "wb") as f:
            ...     f.write(manager.salt)
            >>>
            >>> # Later, to decrypt:
            >>> with open("salt.bin", "rb") as f:
            ...     saved_salt = f.read()
            >>> manager2 = EncryptionManager.from_password("my-password", salt=saved_salt)
            >>> decrypted = manager2.decrypt(encrypted)
        """
        if salt is None:
            salt = os.urandom(16)

        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=32,
            salt=salt,
            iterations=iterations,
        )
        key = base64.urlsafe_b64encode(kdf.derive(password.encode()))

        return cls(key, salt=salt)

    def encrypt(self, data: bytes) -> bytes:
        """Encrypt data.

        Args:
            data: Plaintext data to encrypt

        Returns:
            Encrypted data (includes timestamp and IV)

        Example:
            >>> encrypted = manager.encrypt(b"sensitive data")
        """
        return self._fernet.encrypt(data)

    def decrypt(self, encrypted_data: bytes, ttl: Optional[int] = None) -> bytes:
        """Decrypt data with optional TTL validation for replay protection.

        Args:
            encrypted_data: Encrypted data to decrypt
            ttl: Time-to-live in seconds. If provided, decryption fails if data
                 is older than TTL. This prevents replay attacks with old encrypted data.

        Returns:
            Decrypted plaintext data

        Raises:
            InvalidToken: If data is corrupted, key is wrong, or TTL expired

        Example:
            >>> decrypted = manager.decrypt(encrypted_data)
            >>>
            >>> # With TTL (reject data older than 1 hour)
            >>> decrypted = manager.decrypt(encrypted_data, ttl=3600)

        Security Note:
            Use TTL to prevent replay attacks. For example, if decrypting configuration
            files, use TTL to ensure old encrypted configs cannot be restored maliciously.
        """
        try:
            if ttl is not None:
                return self._fernet.decrypt(encrypted_data, ttl=ttl)
            else:
                return self._fernet.decrypt(encrypted_data)
        except InvalidToken:
            logger.error("Failed to decrypt data: invalid token, corrupted data, or TTL expired")
            raise

    def encrypt_str(self, text: str) -> str:
        """Encrypt a string and return base64-encoded result.

        Args:
            text: Plaintext string to encrypt

        Returns:
            Base64-encoded encrypted string
        """
        encrypted = self.encrypt(text.encode("utf-8"))
        return base64.b64encode(encrypted).decode("ascii")

    def decrypt_str(self, encrypted_text: str, ttl: Optional[int] = None) -> str:
        """Decrypt a base64-encoded encrypted string with optional TTL.

        Args:
            encrypted_text: Base64-encoded encrypted string
            ttl: Time-to-live in seconds (optional, for replay protection)

        Returns:
            Decrypted plaintext string

        Raises:
            InvalidToken: If data is corrupted, key is wrong, or TTL expired
        """
        encrypted = base64.b64decode(encrypted_text.encode("ascii"))
        decrypted = self.decrypt(encrypted, ttl=ttl)
        return decrypted.decode("utf-8")


def get_encryption_key_from_env() -> Optional[bytes]:
    """Get encryption key from environment variable.

    Looks for MCA_SDK_ENCRYPTION_KEY environment variable.

    Returns:
        Encryption key if found, None otherwise

    Example:
        >>> # Set environment variable first:
        >>> # export MCA_SDK_ENCRYPTION_KEY="<base64-encoded-key>"
        >>> key = get_encryption_key_from_env()
        >>> if key:
        ...     manager = EncryptionManager(key)
    """
    key_str = os.environ.get("MCA_SDK_ENCRYPTION_KEY")
    # Check for both None and empty string (empty string is invalid)
    if key_str and key_str.strip():
        try:
            return key_str.encode("ascii")
        except Exception as e:
            logger.warning(f"Invalid encryption key in environment: {e}")
    return None


class QueueEncryption:
    """Encrypts/decrypts queue data at rest with rotation support.

    Uses Fernet (AES-128-CBC + HMAC) with MultiFernet for key rotation.
    The first key is used for encryption. All keys can be used for decryption.
    This enables key rotation: add new key as first element, keep old keys.

    **SECURITY NOTES:**
    - **Encryption Algorithm:** Fernet uses AES-128-CBC (not AES-256) with HMAC for authenticated encryption.
      AES-128 is approved by NIST and provides >112-bit security strength, meeting HIPAA Technical Safeguards.
      The 256-bit Fernet key includes encryption key + HMAC key. For organizations requiring AES-256,
      consider alternative implementations using cryptography.hazmat layer.

    - **Key Storage:** NEVER use environment variables for production keys. Use GCP Secret Manager
      (Story 3.9) or equivalent secure key management system. Environment variables expose keys to:
      process listings, core dumps, container logs, and child processes.

    - **Key Rotation:** Pruned keys are securely zeroed from memory to prevent recovery from
      memory dumps or swap files. Audit logging provided for HIPAA compliance.

    - **Authenticated Encryption:** Fernet provides built-in HMAC authentication. No separate
      integrity checks needed - tampering is automatically detected during decryption.

    Examples:
        >>> # Basic usage (DEV/TEST ONLY - use Secret Manager for production)
        >>> key = Fernet.generate_key()
        >>> encryption = QueueEncryption(encryption_keys=[key])
        >>> encrypted = encryption.encrypt(b"sensitive data")
        >>> decrypted = encryption.decrypt(encrypted)

        >>> # Production usage with Secret Manager (requires Story 3.9)
        >>> encryption = QueueEncryption.from_secret_manager(
        ...     "secret://my-project/queue-key/latest"
        ... )

        >>> # Key rotation with secure deletion and audit trail
        >>> old_key = Fernet.generate_key()
        >>> encryption = QueueEncryption(encryption_keys=[old_key])
        >>> encrypted_old = encryption.encrypt(b"old data")
        >>>
        >>> new_key = Fernet.generate_key()
        >>> encryption.rotate_keys(new_key)
        >>> # Can still decrypt old data, new encryptions use new key
        >>> # Old keys are securely zeroed when pruned
        >>> decrypted = encryption.decrypt(encrypted_old)
    """

    def __init__(
        self, encryption_keys: List[bytes], max_keys: int = 3, default_ttl: Optional[int] = None
    ):
        """Initialize encryption with one or more keys.

        Args:
            encryption_keys: List of 32-byte encryption keys (URL-safe base64-encoded).
                            MUST NOT be empty - raises ConfigurationError if empty.
            max_keys: Maximum number of keys to keep in rotation list (default: 3).
                     Oldest keys beyond this limit are pruned during rotation.
            default_ttl: Default time-to-live in seconds for encrypted data (optional).
                        If set, decryption will fail for data older than TTL.

        Raises:
            ConfigurationError: If encryption_keys is empty or any key is invalid
        """
        if not encryption_keys:
            raise ConfigurationError(
                "No encryption keys provided. QueueEncryption requires at least one valid key. "
                "Generate a key with: from cryptography.fernet import Fernet; key = Fernet.generate_key(). "
                "For production, store keys in GCP Secret Manager or secure key management system."
            )

        if max_keys < 1:
            raise ConfigurationError(f"max_keys must be >= 1, got {max_keys}")

        # Validate all keys before accepting them
        self._validate_keys(encryption_keys)

        self._keys = encryption_keys
        self._max_keys = max_keys
        self._default_ttl = default_ttl
        self._lock = threading.Lock()  # Thread safety for rotation

        # MultiFernet supports rotation: encrypts with first key, decrypts with any
        self._fernet = MultiFernet([Fernet(key) for key in self._keys])

        logger.info(
            f"QueueEncryption initialized with {len(self._keys)} key(s), "
            f"max_keys={max_keys}, default_ttl={default_ttl or 'None'}"
        )

    @staticmethod
    def _validate_keys(keys: List[bytes]) -> None:
        """Validate that all keys are valid Fernet keys.

        Args:
            keys: List of keys to validate

        Raises:
            ConfigurationError: If any key is invalid
        """
        for i, key in enumerate(keys):
            # Fernet keys must be 44 bytes (base64-encoded 32-byte key)
            if not isinstance(key, bytes):
                raise ConfigurationError(
                    f"Key #{i+1} is not bytes type. "
                    f"Expected base64-encoded bytes, got {type(key).__name__}"
                )

            if len(key) != 44:
                raise ConfigurationError(
                    f"Key #{i+1} has invalid length: {len(key)} bytes. "
                    f"Fernet keys must be exactly 44 bytes (base64-encoded 32-byte key). "
                    f"Generate valid key with: from cryptography.fernet import Fernet; "
                    f"key = Fernet.generate_key()"
                )

            # Try to create Fernet instance to validate format
            try:
                Fernet(key)
            except Exception as e:
                raise ConfigurationError(
                    f"Key #{i+1} is invalid: {e}. " f"Must be a valid base64-encoded Fernet key."
                )

    @staticmethod
    def _classify_decryption_error(error: InvalidToken, ttl: Optional[int]) -> str:
        """Classify decryption error type for internal debugging.

        Args:
            error: InvalidToken exception from Fernet
            ttl: TTL value used (None if no TTL)

        Returns:
            Error classification string for logging

        Note:
            This is for INTERNAL logging only. External error messages must
            remain generic to prevent information leakage.

        Implementation Reality:
            The cryptography library's InvalidToken exception typically has NO message
            (empty string) to avoid side-channel leaks. We cannot distinguish between
            "wrong key", "corrupted data", or "TTL expired" based on exception content.

            We can ONLY use context (whether TTL was configured) for classification.
        """
        # InvalidToken message is typically empty - cannot parse it
        # Use context-based classification only

        if ttl is not None:
            # TTL was configured - failure could be expired timestamp OR wrong key
            return "InvalidToken_WithTTL"
        else:
            # No TTL - failure is wrong key OR corrupted data
            # (Fernet validates HMAC before decryption, so usually wrong key)
            return "InvalidToken_NoTTL"

    @staticmethod
    def _key_fingerprint(key: bytes) -> str:
        """Generate SHA-256 fingerprint of key for audit logging.

        Args:
            key: Encryption key

        Returns:
            Hex-encoded SHA-256 hash (first 16 chars for readability)
        """
        return hashlib.sha256(key).hexdigest()[:16]

    def encrypt(self, data: bytes) -> bytes:
        """Encrypt data using the primary (first) key.

        Uses Fernet authenticated encryption which provides:
        - Confidentiality (AES-128-CBC)
        - Authenticity (HMAC-SHA256)
        - Timestamp (for TTL validation)
        - IV (unique per encryption)

        Args:
            data: Plaintext data to encrypt

        Returns:
            Encrypted Fernet token (self-contained, authenticated)

        Example:
            >>> encrypted = encryption.encrypt(b"sensitive data")

        Note:
            Fernet tokens are self-describing and tamper-proof. Any modification
            to the ciphertext will cause decryption to fail with InvalidToken.

        Performance Optimization:
            Lock is acquired only to copy the MultiFernet reference (fast pointer copy),
            then released before the expensive encryption operation. This allows
            concurrent encryption operations for high throughput.

            Python GC Safety: Even if key rotation prunes keys from self._keys and
            creates a new MultiFernet object, the local fernet_instance reference
            keeps the old MultiFernet alive until encryption completes. Python's
            reference counting prevents the object from being garbage collected
            while we hold a reference.

            Therefore, the "key pruning race" cannot cause data loss - the keys
            remain valid in the referenced MultiFernet instance.
        """
        # Acquire lock only to copy MultiFernet reference (fast pointer copy)
        with self._lock:
            fernet_instance = self._fernet

        # Perform encryption WITHOUT holding lock (allows concurrency)
        # Python GC keeps fernet_instance alive until this function returns
        return fernet_instance.encrypt(data)

    def decrypt(self, encrypted_data: bytes, ttl: Optional[int] = None) -> bytes:
        """Decrypt data using any valid key with optional TTL validation.

        Fernet provides authenticated encryption - any tampering or corruption
        will cause decryption to fail with InvalidToken.

        Args:
            encrypted_data: Encrypted Fernet token to decrypt
            ttl: Time-to-live in seconds. If provided, decryption fails if data
                 is older than TTL. Overrides default_ttl if set.

        Returns:
            Decrypted plaintext data

        Raises:
            InvalidToken: If data is corrupted, key is wrong, or TTL expired

        Example:
            >>> # Decrypt with default TTL
            >>> decrypted = encryption.decrypt(encrypted_data)
            >>>
            >>> # Decrypt with custom TTL (30 days)
            >>> decrypted = encryption.decrypt(encrypted_data, ttl=2592000)

        Security Note:
            Error messages are intentionally generic to prevent information
            leakage that could assist cryptographic attacks. Detailed errors
            are logged internally for debugging.

        Performance Note:
            Lock is acquired only to read MultiFernet instance, then released
            before expensive decryption operation. This allows concurrent
            decryption operations for better throughput under load.

        Race Condition Handling:
            If key rotation occurs during decryption, the method implements
            double-check locking with retry:
            1. Thread A reads self._fernet (keys=[v1])
            2. Thread B rotates keys (keys=[v2, v1])
            3. Thread A decrypts - fails with InvalidToken
            4. Thread A re-acquires lock, gets updated keys=[v2, v1]
            5. Thread A retries decryption with new key list
            6. Success - no data loss

            This prevents data loss from race conditions while maintaining
            high concurrency performance (only locks on retry, which is rare).
        """
        # Acquire lock briefly to get current MultiFernet instance
        with self._lock:
            fernet_instance = self._fernet
            effective_ttl = ttl if ttl is not None else self._default_ttl

        # Perform decryption WITHOUT holding lock (allows concurrency)
        try:
            if effective_ttl is not None:
                # Decrypt with TTL validation
                return fernet_instance.decrypt(encrypted_data, ttl=effective_ttl)
            else:
                # Decrypt without TTL
                return fernet_instance.decrypt(encrypted_data)
        except InvalidToken:
            # RACE CONDITION HANDLING: Key rotation may have occurred during decryption
            # Implement double-check locking: re-acquire lock and retry once with updated keys
            with self._lock:
                updated_fernet = self._fernet

            # Only retry if keys actually changed (avoids infinite loop on real errors)
            if updated_fernet is not fernet_instance:
                logger.debug(
                    "Decryption failed with stale keys, retrying with updated key list "
                    "(likely due to concurrent key rotation)"
                )
                try:
                    if effective_ttl is not None:
                        return updated_fernet.decrypt(encrypted_data, ttl=effective_ttl)
                    else:
                        return updated_fernet.decrypt(encrypted_data)
                except InvalidToken as retry_error:
                    # Retry also failed - use retry error for classification (more recent)
                    logger.debug("Retry with updated keys also failed - data truly invalid")
                    e = retry_error  # Use retry exception for error classification

            # Classify error type for internal debugging
            error_class = self._classify_decryption_error(e, effective_ttl)

            # Generic error message to prevent information leakage (external)
            # Detailed classification in logs (internal only)
            logger.error(
                "Decryption failed. Possible causes: wrong key, corrupted data, or TTL expired.",
                extra={
                    "error_class": error_class,
                    "ttl_configured": effective_ttl,
                    "data_size": len(encrypted_data),
                },
            )
            # CRITICAL: Use 'raise e' not bare 'raise'
            # Bare 'raise' re-raises original exception from active context,
            # ignoring our e=retry_error assignment above
            raise e

    def rotate_keys(self, new_key: bytes) -> None:
        """Rotate to a new primary key with audit trail and key limit enforcement.

        The new key is inserted at the beginning of the key list and becomes
        the primary key for all new encryptions. Old keys remain available
        for decrypting existing data, up to max_keys limit.

        **IMPORTANT:** This operation generates audit logs with key fingerprints
        for HIPAA compliance. Ensure audit logs are captured and stored securely.

        Args:
            new_key: New primary encryption key (44-byte base64-encoded Fernet key)

        Raises:
            ConfigurationError: If new_key is invalid

        Example:
            >>> new_key = Fernet.generate_key()
            >>> encryption.rotate_keys(new_key)
            >>> # New encryptions use new_key, old data still decryptable
        """
        with self._lock:
            # Validate new key
            self._validate_keys([new_key])

            # Generate fingerprints for audit trail
            new_key_fp = self._key_fingerprint(new_key)
            old_key_fps = [self._key_fingerprint(k) for k in self._keys[:3]]  # Log first 3

            # Insert new key at beginning (becomes primary)
            self._keys.insert(0, new_key)

            # Enforce max_keys limit by pruning oldest keys
            pruned_keys = []
            if len(self._keys) > self._max_keys:
                pruned_keys = self._keys[self._max_keys :]
                pruned_fps = [self._key_fingerprint(k) for k in pruned_keys]

                # Remove references to pruned keys
                # Note: Python bytes are immutable - cannot be zeroed from memory
                # Keys will remain in memory until garbage collection
                self._keys = self._keys[: self._max_keys]

                logger.warning(
                    f"Key rotation exceeded max_keys limit ({self._max_keys}). "
                    f"Pruned {len(pruned_keys)} oldest key(s). "
                    f"Pruned key fingerprints: {pruned_fps}. "
                    f"Note: Pruned keys remain in memory until garbage collection (Python limitation). "
                    f"Data encrypted with pruned keys will no longer be decryptable."
                )

            # Recreate MultiFernet with updated key list
            self._fernet = MultiFernet([Fernet(key) for key in self._keys])

            # AUDIT LOG: Key rotation event with fingerprints
            logger.info(
                f"AUDIT: Encryption key rotated successfully. "
                f"Timestamp: {time.time()}, "
                f"New key fingerprint: {new_key_fp}, "
                f"Previous key fingerprints: {old_key_fps}, "
                f"Total active keys: {len(self._keys)}/{self._max_keys}, "
                f"Keys pruned: {len(pruned_keys)}"
            )

    @classmethod
    def _unsafe_from_secret_manager_single_key(cls, secret_url: str) -> "QueueEncryption":
        """UNSAFE: Create encryption from Secret Manager (single key only - causes data loss).

        **WARNING: DO NOT USE IN PRODUCTION**

        This method fetches ONLY the latest key version. After key rotation,
        your application WILL LOSE DATA encrypted with old keys on restart.

        Args:
            secret_url: Secret Manager URL (secret://project/secret/version)

        Returns:
            QueueEncryption instance with SINGLE key (data loss risk)

        Raises:
            ImportError: If secret manager module not available (Story 3.9)

        **GUARANTEED DATA LOSS SCENARIO:**
            1. App starts, fetches key v1, encrypts queue data
            2. Key rotation occurs in Secret Manager (v1 -> v2)
            3. App restarts, fetches key v2 only
            4. Cannot decrypt queue data encrypted with v1
            5. Queue data lost permanently

        **Production Solution:**
            Manually fetch and manage multiple key versions:

            >>> # Production-safe approach (fetch multiple versions)
            >>> from your_secret_manager import fetch_key_versions
            >>> keys = fetch_key_versions("queue-key", versions=["v2", "v1"])  # newest first
            >>> encryption = QueueEncryption(encryption_keys=keys)

        **Method Naming:**
            Named with leading underscore and explicit "unsafe" + "single_key" to prevent
            accidental production use. Story 3.9 will provide a safe multi-version fetcher.
        """
        try:
            from .secret_manager import SecretManagerClient
        except ImportError:
            raise ImportError(
                "Secret Manager integration not available. " "Requires Story 3.9 implementation."
            )

        sm_client = SecretManagerClient()
        # Fetch latest key (already base64-encoded 44-byte Fernet key)
        key = sm_client.fetch_secret(secret_url)

        # Validate it's bytes (Secret Manager might return string)
        if isinstance(key, str):
            key = key.encode("ascii")

        return cls(encryption_keys=[key])
