import base64
import functools

from django.db import models
from django.conf import settings
from django.core.exceptions import ImproperlyConfigured
from django.utils.encoding import force_bytes
from cryptography.fernet import Fernet, InvalidToken
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC


@functools.lru_cache(maxsize=128)
def get_legacy_fernet_key(secret_key: str) -> bytes:
    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA256(),
        length=32,
        salt=secret_key.encode("utf-8"),
        iterations=100000,
    )
    return base64.urlsafe_b64encode(kdf.derive(secret_key.encode("utf-8")))


class RotatingFernet:
    """
    A Fernet-based encryption service that supports key rotation.
    
    This class manages encryption/decryption with support for multiple keys,
    enabling seamless key rotation by trying fallback keys during decryption.
    """

    def __init__(self, salt: str, secret_key: str, fallback_secret_keys: list[str] = None):
        self.salt = salt
        self.secret_key = secret_key
        self.fallback_secret_keys = fallback_secret_keys or []
        # Cache Fernet instances to avoid expensive KDF recalculation
        self._fernet_cache: dict[str, Fernet] = {}

    def _get_fernet_key(self, secret_key: str) -> bytes:
        """Generate and cache the Fernet key for a given secret key."""
        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=32,
            salt=force_bytes(self.salt),
            iterations=200000,
        )
        return base64.urlsafe_b64encode(kdf.derive(secret_key.encode("utf-8")))

    def _get_fernet(self, secret_key: str) -> Fernet:
        """Get or create a Fernet instance for the given secret key."""
        if secret_key not in self._fernet_cache:
            fernet_key = self._get_fernet_key(secret_key)
            self._fernet_cache[secret_key] = Fernet(fernet_key)
        return self._fernet_cache[secret_key]

    def encrypt(self, raw_value: bytes) -> bytes:
        """Encrypt a value using the current secret key."""
        return self._get_fernet(self.secret_key).encrypt(raw_value)

    def decrypt(self, encrypted_value: bytes) -> bytes:
        """
        Decrypt a value, trying the current key and all fallback keys.
        
        Raises InvalidToken if decryption fails with all available keys.
        """
        for secret_key in [self.secret_key, *self.fallback_secret_keys]:
            try:
                fernet = self._get_fernet(secret_key)
                return fernet.decrypt(encrypted_value)
            except InvalidToken:
                continue
        raise InvalidToken("Failed to decrypt value with any available key")


@functools.lru_cache(maxsize=1)
def get_default_fernet() -> RotatingFernet:
    return RotatingFernet(
        salt=settings.DATABASE_FIELD_ENCRYPTION_SALT,
        secret_key=settings.DATABASE_FIELD_ENCRYPTION_SECRET_KEY,
        fallback_secret_keys=settings.DATABASE_FIELD_ENCRYPTION_FALLBACK_SECRET_KEYS,
    )


class EncryptedTextField(models.TextField):
    def __init__(self, *args, **kwargs):
        try:
            self._legacy_fernet = Fernet(get_legacy_fernet_key(settings.SECRET_KEY))
            self._fernet = get_default_fernet()
        except Exception as e:
            raise ImproperlyConfigured("Failed to generate Fernet key") from e
        super().__init__(*args, **kwargs)

    def from_db_value(self, value, expression, connection):
        if value is None:
            return value
        try:
            # Ensure value is bytes for decryption
            value_bytes = force_bytes(value)
            try:
                decrypted_bytes = self._fernet.decrypt(value_bytes)
            except InvalidToken:
                decrypted_bytes = self._legacy_fernet.decrypt(value_bytes)
            return decrypted_bytes.decode("utf-8")
        except InvalidToken:
            # Handle cases where decryption fails (e.g., invalid token, wrong key)
            # Instead of returning the potentially sensitive raw value, raise an error.
            # Alternatively, log the error and return None or a default value if appropriate.
            raise ValueError(
                "Failed to decrypt database value. Check key or data integrity."
            )
        except Exception as e:
            # Catch other potential errors during the process (e.g., decoding)
            # Re-raising or logging might be appropriate here too.
            # For now, let's re-raise to make unexpected issues visible.
            raise ValueError(
                f"An unexpected error occurred during decryption/decoding: {e}"
            ) from e

    def to_python(self, value):
        if value is None:
            return value
        # This method converts input values (e.g., from forms) to Python objects.
        # It should not perform decryption. Decryption happens in from_db_value.
        # For a TextField, we usually just ensure it's a string.
        if isinstance(value, bytes):
            # If bytes are somehow passed (e.g., from db raw query), decode them
            # This might be redundant if from_db_value is always called first
            try:
                return value.decode("utf-8")
            except UnicodeDecodeError:
                # Handle cases where bytes are not valid UTF-8, perhaps raise error
                # or return as is depending on requirements.
                raise ValueError(
                    "Invalid byte sequence for UTF-8 decoding in to_python"
                )
        elif not isinstance(value, str):
            # Convert non-string types if necessary/possible
            return str(value)

        # If it's already a string or None (handled above), return as is.
        return value

    def get_prep_value(self, value):
        if value is None:
            return value
        try:
            return self._fernet.encrypt(force_bytes(value)).decode("utf-8")
        except Exception as e:
            raise ValueError(f"Failed to encrypt value: {e}")