"""Symmetric encryption for wallet keys with versioning."""
import base64
import hashlib
from cryptography.fernet import Fernet
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC


class KeyEncryption:
    """
    Versioned encryption for wallet keys.
    Supports key rotation by maintaining multiple key versions.
    """
    CURRENT_VERSION = 1

    def __init__(self, master_keys: dict):
        """
        Initialize with versioned master keys.

        Args:
            master_keys: Dict mapping version numbers to master key strings.
                         Example: {1: 'key_v1', 2: 'key_v2'}
        """
        self.fernets = {}
        for version, key in master_keys.items():
            # Derive unique salt per version (deterministic but version-specific)
            salt = hashlib.sha256(f"brighter_wallet_v{version}".encode()).digest()[:16]
            kdf = PBKDF2HMAC(
                algorithm=hashes.SHA256(),
                length=32,
                salt=salt,
                iterations=100000,
            )
            derived_key = base64.urlsafe_b64encode(kdf.derive(key.encode()))
            self.fernets[int(version)] = Fernet(derived_key)

    def encrypt(self, data: str, version: int = None) -> str:
        """
        Encrypt data with current (or specified) key version.

        Args:
            data: Plaintext string to encrypt.
            version: Key version to use (defaults to CURRENT_VERSION).

        Returns:
            Encrypted data as string.
        """
        version = version or self.CURRENT_VERSION
        if version not in self.fernets:
            raise ValueError(f"Unknown encryption key version: {version}")
        return self.fernets[version].encrypt(data.encode()).decode()

    def decrypt(self, encrypted_data: str, version: int) -> str:
        """
        Decrypt data using the specified key version.

        Args:
            encrypted_data: Encrypted string to decrypt.
            version: Key version that was used to encrypt.

        Returns:
            Decrypted plaintext string.

        Raises:
            ValueError: If the key version is unknown.
        """
        if version not in self.fernets:
            raise ValueError(f"Unknown encryption key version: {version}")
        return self.fernets[version].decrypt(encrypted_data.encode()).decode()

    def re_encrypt(self, encrypted_data: str, old_version: int,
                   new_version: int = None) -> tuple[str, int]:
        """
        Re-encrypt data from old version to new version.
        Used during key rotation.

        Args:
            encrypted_data: Data encrypted with old_version.
            old_version: Version the data is currently encrypted with.
            new_version: Target version (defaults to CURRENT_VERSION).

        Returns:
            Tuple of (new_encrypted_data, new_version).
        """
        new_version = new_version or self.CURRENT_VERSION
        plaintext = self.decrypt(encrypted_data, old_version)
        return self.encrypt(plaintext, new_version), new_version
