"""Envelope-encrypted credential vault (build plan Section 9.3).

Each secret is encrypted with a fresh per-record data key (DEK); the DEK is then
wrapped by a master key (KEK) loaded from ``LEGATE_MASTER_KEY``. Multiple KEK
versions are supported so keys can be rotated: a record records the
``key_version`` its DEK was wrapped with, and :meth:`Vault.rewrap` re-wraps a
record under the current key without touching the ciphertext.

Plaintext only ever exists in memory during encrypt/decrypt; ciphertext and the
wrapped DEK are what live at rest.
"""

from __future__ import annotations

import base64
import json
from dataclasses import dataclass
from functools import lru_cache

from cryptography.fernet import Fernet, InvalidToken

from legate.config import Settings, get_settings


class VaultError(Exception):
    """Raised when the vault is misconfigured or decryption fails."""


@dataclass(frozen=True)
class EncryptedSecret:
    ciphertext: str
    wrapped_dek: str
    key_version: str


def _fernet_from_master(master_key_b64: str) -> Fernet:
    if not master_key_b64:
        raise VaultError("LEGATE_MASTER_KEY is not set; cannot use the credential vault")
    try:
        raw = base64.b64decode(master_key_b64)
    except (ValueError, TypeError) as exc:
        raise VaultError("LEGATE_MASTER_KEY must be base64-encoded") from exc
    if len(raw) != 32:
        raise VaultError("LEGATE_MASTER_KEY must decode to exactly 32 bytes")
    return Fernet(base64.urlsafe_b64encode(raw))


class Vault:
    def __init__(self, keys: dict[str, Fernet], current_version: str) -> None:
        if current_version not in keys:
            raise VaultError(f"current key version {current_version!r} is not in the key set")
        self._keys = keys
        self._current_version = current_version

    @classmethod
    def from_settings(cls, settings: Settings | None = None) -> Vault:
        settings = settings or get_settings()
        return cls(keys={"1": _fernet_from_master(settings.legate_master_key)}, current_version="1")

    def _key(self, version: str) -> Fernet:
        key = self._keys.get(version)
        if key is None:
            raise VaultError(f"no master key available for version {version!r}")
        return key

    def encrypt(self, secrets: dict) -> EncryptedSecret:
        dek = Fernet.generate_key()
        plaintext = json.dumps(secrets, separators=(",", ":")).encode("utf-8")
        ciphertext = Fernet(dek).encrypt(plaintext)
        wrapped = self._key(self._current_version).encrypt(dek)
        return EncryptedSecret(
            ciphertext=ciphertext.decode("ascii"),
            wrapped_dek=wrapped.decode("ascii"),
            key_version=self._current_version,
        )

    def decrypt(self, ciphertext: str, wrapped_dek: str, key_version: str) -> dict:
        try:
            dek = self._key(key_version).decrypt(wrapped_dek.encode("ascii"))
            plaintext = Fernet(dek).decrypt(ciphertext.encode("ascii"))
        except InvalidToken as exc:
            raise VaultError("failed to decrypt credential (wrong key or corrupted data)") from exc
        return json.loads(plaintext)

    def rewrap(self, ciphertext: str, wrapped_dek: str, key_version: str) -> EncryptedSecret:
        """Re-wrap a record's DEK under the current key version (rotation)."""
        try:
            dek = self._key(key_version).decrypt(wrapped_dek.encode("ascii"))
        except InvalidToken as exc:
            raise VaultError("failed to unwrap credential during rotation") from exc
        wrapped = self._key(self._current_version).encrypt(dek)
        return EncryptedSecret(
            ciphertext=ciphertext,
            wrapped_dek=wrapped.decode("ascii"),
            key_version=self._current_version,
        )


@lru_cache
def get_vault() -> Vault:
    return Vault.from_settings()
