import os
import logging
from typing import Optional
from cryptography.fernet import Fernet
from atomsql.orm import Model, fields

logger = logging.getLogger(__name__)


class KeyStoreModel(Model):
    """
    Stores encrypted data keys (wrapped by the Master Key).
    Table: atom_keystore
    """

    key_identifier = fields.CharField(max_length=255, unique=True)
    wrapped_key = fields.TextField()  # The data key encrypted by the Master Key
    created_at = fields.DateTimeField()
    # If the key is shredded, we delete the row.


class KeyManager:
    """
    Manages generation, retrieval, and destruction of encryption keys.
    """

    _master_cipher: Optional[Fernet] = None

    @classmethod
    def get_master_cipher(cls) -> Fernet:
        """Get the cipher for the Master Key from env vars."""
        if cls._master_cipher is None:
            key = os.environ.get("ATOMINGEST_MASTER_KEY")
            if not key:
                # Fallback to legacy key if master not set, for backward compat
                key = os.environ.get("ATOMINGEST_ENCRYPTION_KEY")

            if not key:
                raise ValueError(
                    "ATOMINGEST_MASTER_KEY environment variable is not set. "
                    "Crypto-shredding requires a master key."
                )
            try:
                cls._master_cipher = Fernet(key)
            except Exception as e:
                raise ValueError(f"Invalid Master Key: {e}")
        return cls._master_cipher

    @classmethod
    def get_key(cls, key_identifier: str) -> bytes:
        """
        Retrieve the raw data key for a specific identifier.
        If it doesn't exist, create it.
        """
        master = cls.get_master_cipher()

        # Try to find existing key
        # Note: This assumes KeyStoreModel is registered with the active DB connection
        try:
            keystore_entry = list(
                KeyStoreModel.objects().filter(key_identifier=key_identifier)
            )
        except Exception as e:
            logger.error(f"Database error looking up key '{key_identifier}': {e}")
            raise

        if keystore_entry:
            # Decrypt the wrapped key — let decryption errors propagate
            # so corrupted keys are not silently replaced
            decrypted_key: bytes = master.decrypt(
                keystore_entry[0].wrapped_key.encode()
            )
            return decrypted_key
        else:
            # Key doesn't exist, create new
            return cls.create_key(key_identifier)

    @classmethod
    def create_key(cls, key_identifier: str) -> bytes:
        """Generate and store a new data key."""
        master = cls.get_master_cipher()
        new_key: bytes = Fernet.generate_key()
        wrapped_key = master.encrypt(new_key).decode("utf-8")

        try:
            KeyStoreModel.create(key_identifier=key_identifier, wrapped_key=wrapped_key)
            return new_key
        except Exception as e:
            # Handle race condition or DB error
            logger.error(f"Failed to create key for {key_identifier}: {e}")
            raise

    @classmethod
    def shred_key(cls, key_identifier: str) -> bool:
        """
        Crypto-Shredding: Permanently delete the key for an identifier.
        This renders all data encrypted with this key unreadable.
        """
        try:
            entries = list(
                KeyStoreModel.objects().filter(key_identifier=key_identifier)
            )
            if not entries:
                logger.warning(f"Key '{key_identifier}' does not exist for shredding.")
                return False

            entries[0].delete()
            logger.warning(
                f"CRYPTO-SHREDDING: Key for '{key_identifier}' has been destroyed."
            )
            return True
        except Exception as e:
            logger.error(f"Failed to shred key for {key_identifier}: {e}")
            return False
