"""Encryption providers for secrets at rest."""

from __future__ import annotations

import base64
import os
from typing import Protocol, runtime_checkable

from cryptography.fernet import Fernet
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.hazmat.primitives import hashes


@runtime_checkable
class EncryptionProvider(Protocol):
    """Interface that any encryption backend must satisfy."""

    def encrypt(self, plaintext: bytes) -> bytes: ...
    def decrypt(self, ciphertext: bytes) -> bytes: ...


def derive_key(master_password: str, salt: bytes | None = None, length: int = 32) -> tuple[bytes, bytes]:
    """Derive an encryption key from a master password using PBKDF2.

    Returns (derived_key, salt) so the salt can be stored alongside ciphertext.
    """
    if salt is None:
        salt = os.urandom(16)
    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA256(),
        length=length,
        salt=salt,
        iterations=600_000,
    )
    return kdf.derive(master_password.encode()), salt


class FernetEncryption:
    """Fernet symmetric encryption (AES-128-CBC under the hood)."""

    def __init__(self, key: bytes | None = None, master_password: str | None = None) -> None:
        if key is not None:
            self._fernet = Fernet(key)
        elif master_password is not None:
            derived, _ = derive_key(master_password, salt=b"willian-secrets-fernet", length=32)
            fernet_key = base64.urlsafe_b64encode(derived)
            self._fernet = Fernet(fernet_key)
        else:
            raise ValueError("Provide either 'key' (url-safe base64) or 'master_password'")

    def encrypt(self, plaintext: bytes) -> bytes:
        return self._fernet.encrypt(plaintext)

    def decrypt(self, ciphertext: bytes) -> bytes:
        return self._fernet.decrypt(ciphertext)


class AESGCMEncryption:
    """AES-256-GCM authenticated encryption."""

    _NONCE_SIZE = 12  # 96-bit nonce recommended for GCM

    def __init__(self, key: bytes | None = None, master_password: str | None = None) -> None:
        if key is not None:
            if len(key) != 32:
                raise ValueError("AES-256-GCM requires a 32-byte key")
            self._key = key
        elif master_password is not None:
            self._key, _ = derive_key(master_password, salt=b"willian-secrets-aesgcm", length=32)
        else:
            raise ValueError("Provide either 'key' (32 bytes) or 'master_password'")
        self._aesgcm = AESGCM(self._key)

    def encrypt(self, plaintext: bytes) -> bytes:
        nonce = os.urandom(self._NONCE_SIZE)
        ciphertext = self._aesgcm.encrypt(nonce, plaintext, None)
        return nonce + ciphertext  # prepend nonce for decryption

    def decrypt(self, ciphertext: bytes) -> bytes:
        nonce = ciphertext[: self._NONCE_SIZE]
        actual_ct = ciphertext[self._NONCE_SIZE :]
        return self._aesgcm.decrypt(nonce, actual_ct, None)
