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


class KeyEncryption:
    def __init__(self, encryption_key: str | None = None):
        self._fernet = self._create_fernet(encryption_key)

    def _create_fernet(self, key: str | None) -> Fernet:
        if not key:
            key = os.getenv("API_KEY_ENCRYPTION_KEY", "")

        if not key:
            return Fernet(Fernet.generate_key())

        try:
            return Fernet(key.encode())
        except Exception:
            kdf = PBKDF2HMAC(
                algorithm=hashes.SHA256(),
                length=32,
                salt=b"geo-platform-salt",
                iterations=100000,
            )
            derived_key = base64.urlsafe_b64encode(kdf.derive(key.encode()))
            return Fernet(derived_key)

    def encrypt(self, plaintext: str) -> str:
        if not plaintext:
            raise ValueError("Cannot encrypt empty string")
        return self._fernet.encrypt(plaintext.encode()).decode()

    def decrypt(self, ciphertext: str) -> str:
        if not ciphertext:
            raise ValueError("Cannot decrypt empty string")
        try:
            return self._fernet.decrypt(ciphertext.encode()).decode()
        except InvalidToken:
            raise ValueError("Invalid ciphertext or key")


_key_encryption: KeyEncryption | None = None


def get_key_encryption() -> KeyEncryption:
    global _key_encryption
    if _key_encryption is None:
        _key_encryption = KeyEncryption()
    return _key_encryption


def reset_key_encryption() -> None:
    global _key_encryption
    _key_encryption = None