import os
import base64
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.backends import default_backend

class EncryptionManager:
    """
    A utility class to handle secure encryption and decryption of strings
    using a client secret.

    This implementation uses AES-GCM for authenticated encryption, ensuring
    both confidentiality and integrity. The encryption key is derived from
    the client secret using PBKDF2.
    """
    # --- Constants for cryptographic parameters ---
    # AES key size in bytes (32 bytes = 256 bits)
    _AES_KEY_SIZE = 32
    # GCM nonce size in bytes (12 bytes is recommended for GCM)
    _GCM_NONCE_SIZE = 12
    # GCM authentication tag size in bytes (16 bytes is standard)
    _GCM_TAG_SIZE = 16
    # PBKDF2 salt size in bytes (16 bytes is a good minimum)
    _PBKDF2_SALT_SIZE = 16
    # PBKDF2 iterations (higher is more secure but slower)
    _PBKDF2_ITERATIONS = 390_000

    def __init__(self, client_secret: str):
        if not client_secret:
            raise ValueError("Client secret cannot be empty.")
        self._client_secret = client_secret.encode('utf-8')
        self._backend = default_backend()

    def _derive_key(self, salt: bytes) -> bytes:
        """Derives a 256-bit AES key from the client secret using PBKDF2."""
        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=self._AES_KEY_SIZE,
            salt=salt,
            iterations=self._PBKDF2_ITERATIONS,
            backend=self._backend
        )
        return kdf.derive(self._client_secret)

    def encrypt(self, plaintext: str) -> str:
        """
        Encrypts a string using AES-GCM and returns a URL-safe base64 encoded string.

        The output format is: salt + nonce + ciphertext + tag

        Args:
            plaintext (str): The string to encrypt.

        Returns:
            str: A base64 encoded, URL-safe string containing all necessary
                 components for decryption.
        """
        if not isinstance(plaintext, str) or not plaintext:
            raise ValueError("Plaintext must be a non-empty string.")

        # 1. Generate a random salt for key derivation and a nonce for GCM
        salt = os.urandom(self._PBKDF2_SALT_SIZE)
        nonce = os.urandom(self._GCM_NONCE_SIZE)

        # 2. Derive the encryption key from the secret and salt
        key = self._derive_key(salt)

        # 3. Encrypt the plaintext
        encryptor = Cipher(
            algorithms.AES(key),
            modes.GCM(nonce),
            backend=self._backend
        ).encryptor()

        ciphertext = encryptor.update(plaintext.encode('utf-8')) + encryptor.finalize()

        # 4. Combine salt, nonce, ciphertext, and the authentication tag
        encrypted_data = salt + nonce + ciphertext + encryptor.tag

        # 5. Return as a URL-safe base64 encoded string for easy storage/transport
        return base64.urlsafe_b64encode(encrypted_data).decode('utf-8')

    def decrypt(self, encrypted_text: str) -> str:
        """
        Decrypts a base64 encoded string that was encrypted by the `encrypt` method.

        Args:
            encrypted_text (str): The base64 encoded encrypted string.

        Returns:
            str: The original decrypted string.

        Raises:
            ValueError: If decryption fails due to tampered data, incorrect secret,
                        or malformed input.
        """
        try:
            encrypted_data = base64.urlsafe_b64decode(encrypted_text)

            # 1. Extract the components from the combined data
            salt = encrypted_data[:self._PBDF2_SALT_SIZE]
            nonce = encrypted_data[self._PBDF2_SALT_SIZE : self._PBDF2_SALT_SIZE + self._GCM_NONCE_SIZE]
            tag_index = len(encrypted_data) - self._GCM_TAG_SIZE
            ciphertext = encrypted_data[self._PBDF2_SALT_SIZE + self._GCM_NONCE_SIZE : tag_index]
            tag = encrypted_data[tag_index:]

            # 2. Derive the key using the extracted salt
            key = self._derive_key(salt)

            # 3. Decrypt and verify the data
            decryptor = Cipher(
                algorithms.AES(key),
                modes.GCM(nonce, tag),
                backend=self._backend
            ).decryptor()

            decrypted_bytes = decryptor.update(ciphertext) + decryptor.finalize()
            return decrypted_bytes.decode('utf-8')

        except Exception as e:
            # Broad exception to catch any cryptographic error (e.g., InvalidTag)
            # or decoding/slicing error, which indicates invalid input.
            raise ValueError("Decryption failed. The data may be tampered, "
                             "corrupt, or encrypted with a different secret.") from e
