import logging
import json
from typing import Any, Optional

from cryptography.fernet import Fernet
from atomsql.orm import fields

from atomingest.security.kms import KeyManager

logger = logging.getLogger(__name__)


class EncryptedFieldMixin:
    """
    Mixin that provides transparent encryption and decryption for AtomSQL Fields.
    Support granular encryption via key_id_field
    """

    def __init__(self, key_id_field: Optional[str] = None, **kwargs):
        self.key_id_field = key_id_field
        super().__init__(**kwargs)

    def _get_cipher(self, instance: Any) -> Fernet:
        """Get the correct cipher.
        if key_id_field is set, use KeyManager to get the key.
        else, use the global Master key (fallback/simple mode)
        """
        if self.key_id_field and instance is not None:
            key_id = getattr(instance, self.key_id_field, None)
            if key_id:
                raw_key = KeyManager.get_key(str(key_id))
                return Fernet(raw_key)
            else:
                logger.warning(
                    f"key_id_field '{self.key_id_field}' is missing on instance. Using Master Key."
                )
        return KeyManager.get_master_cipher()

    def __set__(self, instance: Any, value: Any):
        if value is None:
            getattr(super(), "__set__")(instance, value)
            return
        cipher = self._get_cipher(instance)

        str_val = (
            json.dumps(value) if isinstance(self, fields.JSONField) else str(value)
        )

        try:
            encrypted_value = cipher.encrypt(str_val.encode("utf-8")).decode("utf-8")
            instance.__dict__[self.name] = encrypted_value  # type: ignore[attr-defined]
        except Exception as e:
            logger.error(f"Encryption failed for field '{self.name}': {e}")  # type: ignore[attr-defined]
            raise

    def __get__(self, instance: Any, owner: Any) -> Any:
        encrypted_value = instance.__dict__.get(self.name)  # type: ignore[attr-defined]
        if encrypted_value is None:
            return None
        cipher = self._get_cipher(instance)
        try:
            decrypted_bytes = cipher.decrypt(encrypted_value.encode("utf-8"))
            decrypted_str = decrypted_bytes.decode("utf-8")
            if isinstance(self, fields.JSONField):
                return json.loads(decrypted_str)
            return decrypted_str
        except Exception as e:
            logger.error(f"Decryption failed for field '{self.name}': {e}")  # type: ignore[attr-defined]
            raise


class EncryptedStringField(EncryptedFieldMixin, fields.StringField):
    """
    String field with transparent encryption/decryption.
    Inherits from AtomSQL's StringField and EncryptedFieldMixin.
    """

    pass


class EncryptedCharField(EncryptedFieldMixin, fields.CharField):
    """
    Char field with transparent encryption/decryption.
    Inherits from AtomSQL's CharField and EncryptedFieldMixin.
    """

    pass


class EncryptedTextField(EncryptedFieldMixin, fields.TextField):
    """
    Text field with transparent encryption/decryption.
    Inherits from AtomSQL's TextField and EncryptedFieldMixin.
    """

    pass


class EncryptedJSONField(EncryptedFieldMixin, fields.JSONField):
    """
    JSON field with transparent encryption/decryption.
    Inherits from AtomSQL's JSONField and EncryptedFieldMixin.
    """

    def to_database(self, value: Any) -> Optional[str]:
        if value is None:
            return None
        json_str = json.dumps(value)
        return super().to_database(json_str)  # type: ignore[no-any-return, misc]

    def to_python(self, value: Any) -> Any:
        decrypted_str = super().to_python(value)
        if decrypted_str is None:
            return None
        return json.loads(decrypted_str)


class EncryptedEmailField(EncryptedFieldMixin, fields.EmailField):
    """
    Email field with transparent encryption/decryption.
    Inherits from AtomSQL's EmailField and EncryptedFieldMixin.
    """

    pass
