"""Buffer encryption for offline payloads.

Buffered discovery payloads contain only infrastructure metadata (model names,
endpoint URLs, latency/error aggregates) — never prompt or completion content.
They are still encrypted at rest with a key derived from the tenant's
``AISSURANCE_KEY`` via HKDF-SHA256, so a leaked buffer file is useless without
the tenant secret.

The on-disk format is Fernet (AES-128-CBC + HMAC-SHA256), which authenticates
the ciphertext and embeds a timestamp.
"""
import base64
from typing import Optional

from cryptography.fernet import Fernet, InvalidToken
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.hkdf import HKDF

# Domain-separation label so this derived key can never collide with a key
# derived from the same secret for a different purpose.
_HKDF_INFO = b"aissurance-buffer-v1"


def _fernet(api_key: str) -> Fernet:
    """Derive a Fernet instance from the tenant API key."""
    derived = HKDF(
        algorithm=hashes.SHA256(),
        length=32,
        salt=None,
        info=_HKDF_INFO,
    ).derive(api_key.encode("utf-8"))
    return Fernet(base64.urlsafe_b64encode(derived))


def encrypt(plaintext: bytes, api_key: str) -> bytes:
    """Encrypt ``plaintext`` for at-rest buffer storage."""
    return _fernet(api_key).encrypt(plaintext)


def decrypt(token: bytes, api_key: str) -> Optional[bytes]:
    """Decrypt a buffer file. Returns None if the key is wrong or data is corrupt."""
    try:
        return _fernet(api_key).decrypt(token)
    except (InvalidToken, ValueError):
        return None