"""
Envelope encryption key management for LockStock.

Provides the cryptographic primitives for the envelope encryption layer
(Layer 4) that composes WITH chain-as-key binding to add confidentiality.

The key hierarchy:

    Account Key (operator secret, ONE per organization)
         │
         ├─ KEK = HKDF(account_key, agent_id)    # per-agent, deterministic
         │      └─ wraps/unwraps DEK
         │
         └─ DEK (per-agent, random, stored WRAPPED on LockStock)
                └─ blob_key = HKDF(DEK, stamp_hash, payload_hash)  # per-blob

What each party holds:
    Operator:   Account Key (the ONLY secret they manage)
    Agent:      DEK in memory (unwrapped from server on boot)
    LockStock:  Wrapped DEK (opaque bytes, cannot unwrap without Account Key)
    Auditor:    Receives Account Key from operator through governance process

IMPORTANT — Agent ID is cryptographic identity:
    derive_kek() uses agent_id as HKDF salt. If an agent's ID changes
    (re-registration, migration), the KEK changes and ALL previously
    wrapped DEKs for that agent become unrecoverable. Agent ID is not
    just a label — it is load-bearing for key derivation. Document this
    in operator-facing materials.

Security properties:
    - LockStock cannot decrypt ANY customer data (no Account Key, no DEK)
    - Compromised DEK from one agent does not compromise other agents
      (KEK is per-agent via HKDF with agent_id salt)
    - Per-blob keys are chain-bound: tamper/position/ordering binding preserved
    - Key rotation: new DEK version, old blobs still decryptable with old DEK
"""

import hashlib
import os

from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
from cryptography.hazmat.primitives import hashes


# --- Context strings (NEVER CHANGE — changing these breaks all encrypted data) ---

KEK_CONTEXT = b"lockstock-kek-v1"
BLOB_KEY_CONTEXT_PREFIX = b"lockstock-payload-v1:"
ACCOUNT_KEY_VERIFY_PREFIX = b"lockstock-account-key-verify-v1"


# --- Account Key ---

def generate_account_key() -> bytes:
    """
    Generate a 256-bit Account Key.

    The operator stores this key securely (environment variable, vault, KMS).
    It is the ONLY secret the operator manages across all their agents.
    LockStock never sees this key — only a verification hash.

    Returns:
        32 bytes of cryptographically random data.
    """
    return os.urandom(32)


def account_key_verification_hash(account_key: bytes) -> str:
    """
    One-way hash of the Account Key for server-side verification.

    The server stores this hash to verify the operator "knows" their key
    (e.g., during key management operations). The server CANNOT derive
    the Account Key from this hash.

    Args:
        account_key: 32-byte Account Key

    Returns:
        64-char hex string (SHA-256)
    """
    return hashlib.sha256(
        ACCOUNT_KEY_VERIFY_PREFIX + account_key
    ).hexdigest()


# --- KEK Derivation ---

def derive_kek(account_key: bytes, agent_id: str) -> bytes:
    """
    Derive a per-agent Key Encryption Key from the Account Key.

    Deterministic: same (account_key, agent_id) always produces the same KEK.
    This means a compromised DEK from agent A does NOT give the attacker
    the KEK to unwrap agent B's DEK (different agent_id = different salt
    = different KEK), even under the same Account Key.

    IMPORTANT: agent_id is cryptographic identity here. If the agent's ID
    changes, the KEK changes, and all previously wrapped DEKs become
    unrecoverable. Never rename or reassign agent IDs.

    Args:
        account_key: 32-byte Account Key (operator's secret)
        agent_id: The agent's unique identifier (used as HKDF salt)

    Returns:
        32 bytes suitable for AES-256-GCM key wrapping.
    """
    hkdf = HKDF(
        algorithm=hashes.SHA256(),
        length=32,
        salt=agent_id.encode("utf-8"),
        info=KEK_CONTEXT,
    )
    return hkdf.derive(account_key)


# --- DEK Generation & Wrapping ---

def generate_dek() -> bytes:
    """
    Generate a 256-bit Data Encryption Key.

    One DEK per agent. Wrapped with the KEK and stored on LockStock.
    The plaintext DEK never leaves the agent's process memory.

    Returns:
        32 bytes of cryptographically random data.
    """
    return os.urandom(32)


def wrap_dek(kek: bytes, dek: bytes) -> bytes:
    """
    Wrap a DEK using AES-256-GCM with a random nonce.

    The wrapped output is self-contained: nonce is prepended.
    Each wrapping produces different output (random nonce) even for
    the same (KEK, DEK) pair.

    Format: nonce(12 bytes) + ciphertext(32 bytes) + tag(16 bytes) = 60 bytes

    Args:
        kek: 32-byte Key Encryption Key (from derive_kek)
        dek: 32-byte Data Encryption Key to wrap

    Returns:
        60 bytes: nonce(12) || AES-GCM(kek, nonce, dek)
    """
    nonce = os.urandom(12)
    aesgcm = AESGCM(kek)
    ciphertext_and_tag = aesgcm.encrypt(nonce, dek, None)
    return nonce + ciphertext_and_tag


def unwrap_dek(kek: bytes, wrapped_dek: bytes) -> bytes:
    """
    Unwrap a DEK.

    Splits the nonce (first 12 bytes) from the ciphertext+tag,
    then decrypts with AES-256-GCM.

    Args:
        kek: 32-byte Key Encryption Key (from derive_kek)
        wrapped_dek: 60 bytes from wrap_dek()

    Returns:
        32-byte DEK.

    Raises:
        cryptography.exceptions.InvalidTag: if wrong KEK or tampered data
    """
    nonce = wrapped_dek[:12]
    ciphertext_and_tag = wrapped_dek[12:]
    aesgcm = AESGCM(kek)
    return aesgcm.decrypt(nonce, ciphertext_and_tag, None)


# --- Blob Key Derivation ---

def derive_blob_key(dek: bytes, stamp_hash: str, payload_hash: str) -> bytes:
    """
    Derive a per-blob encryption key by composing envelope encryption
    with chain-as-key binding.

    This is the core innovation: the DEK provides confidentiality (gate),
    the stamp_hash provides position binding (chain), and the payload_hash
    provides content binding. All three are required to produce the key.

    Anyone missing the DEK (including LockStock) cannot derive any blob key,
    even with full chain access. Anyone with the DEK but the wrong chain
    position or wrong content gets a different key — decryption fails.

    HKDF parameters:
        ikm  = DEK                    (confidentiality gate)
        salt = stamp_hash bytes       (chain position binding)
        info = prefix + payload_hash  (content binding + domain separation)

    Args:
        dek: 32-byte Data Encryption Key
        stamp_hash: 64-char hex string of the stamp's hash on chain
        payload_hash: 64-char hex string of SHA-256(canonical_bytes)

    Returns:
        32 bytes suitable for AES-256-GCM.
    """
    hkdf = HKDF(
        algorithm=hashes.SHA256(),
        length=32,
        salt=bytes.fromhex(stamp_hash),
        info=BLOB_KEY_CONTEXT_PREFIX + bytes.fromhex(payload_hash),
    )
    return hkdf.derive(dek)


def derive_blob_nonce(dek: bytes, stamp_hash: str) -> bytes:
    """
    Derive a 96-bit nonce for blob encryption.

    SAFETY ARGUMENT FOR DETERMINISTIC NONCE:
    AES-GCM requires that (key, nonce) pairs are never reused. A deterministic
    nonce is normally dangerous, but is safe here because:

    1. The blob_key is unique per blob — it incorporates payload_hash via
       HKDF info. Two different payloads at the same chain position get
       different keys, so even with the same nonce, the (key, nonce) pair
       is unique.

    2. The stamp_hash is unique per chain position (chain invariant). Two
       blobs at different positions get different nonces AND different keys.

    3. The DEK is included in the nonce derivation. After key rotation,
       even the same stamp_hash produces a different nonce with the new DEK.

    Therefore: different payload → different key. Different position → different
    nonce AND different key. Different DEK → different nonce AND different key.
    No (key, nonce) reuse is possible.

    Args:
        dek: 32-byte Data Encryption Key
        stamp_hash: 64-char hex string of the stamp's hash on chain

    Returns:
        12 bytes (96-bit nonce for AES-GCM).
    """
    return hashlib.sha256(
        dek + bytes.fromhex(stamp_hash)
    ).digest()[:12]
