"""MFP Storage Encryption — per-cell AES-256-GCM for at-rest columns.

Encrypts sensitive BLOB columns (local_state, ratchet_state, shared_prng_seed)
before writing to SQLite. Decrypts on read.

Security audit (Phase 6, task 6.3):
- Fixed: nonce derivation changed from deterministic to random (os.urandom).
  Deterministic nonces revealed whether plaintext changed between writes.
- AAD binding: length-prefixed concatenation prevents collision.
- Key derivation: HMAC-SHA-256 KDF chain (master → storage → cell) verified.

Maps to: impl/I-13_schema.md §6
"""

from __future__ import annotations

import os
import struct

from mfp.core.primitives import aes_256_gcm_decrypt, aes_256_gcm_encrypt, hmac_sha256
from mfp.core.types import StateValue


def derive_storage_key(master_key: bytes, runtime_id: bytes) -> StateValue:
    """Derive the storage-wide encryption key from master key + runtime ID.

    storage_key = HMAC-SHA-256(master_key, "mfp-storage-key" || runtime_id)
    """
    return hmac_sha256(master_key, b"mfp-storage-key" + runtime_id)


def derive_cell_key(
    storage_key: StateValue,
    table_name: bytes,
    column_name: bytes,
    row_id: bytes,
) -> StateValue:
    """Derive a per-cell encryption key.

    cell_key = HMAC-SHA-256(storage_key, len(table) || table || len(column) || column || row_id)

    Length-prefixed to prevent collisions (e.g., table="ab" column="c" vs table="a" column="bc").
    """
    msg = struct.pack(">H", len(table_name)) + table_name + struct.pack(">H", len(column_name)) + column_name + row_id
    return hmac_sha256(storage_key.data, msg)


def _make_aad(row_id: bytes, column_name: bytes) -> bytes:
    """Build length-prefixed AAD for cell encryption."""
    return struct.pack(">H", len(row_id)) + row_id + struct.pack(">H", len(column_name)) + column_name


def encrypt_cell(
    storage_key: StateValue,
    table_name: bytes,
    column_name: bytes,
    row_id: bytes,
    plaintext: bytes,
) -> bytes:
    """Encrypt a cell value. Returns random_nonce || ciphertext || tag."""
    cell_key = derive_cell_key(storage_key, table_name, column_name, row_id)
    nonce = os.urandom(12)  # Random nonce — prevents ciphertext comparison attacks
    aad = _make_aad(row_id, column_name)
    ciphertext = aes_256_gcm_encrypt(cell_key.data, nonce, plaintext, aad)
    return nonce + ciphertext


def decrypt_cell(
    storage_key: StateValue,
    table_name: bytes,
    column_name: bytes,
    row_id: bytes,
    encrypted: bytes,
) -> bytes | None:
    """Decrypt a cell value. Returns None on integrity failure."""
    cell_key = derive_cell_key(storage_key, table_name, column_name, row_id)
    nonce = encrypted[:12]
    ciphertext = encrypted[12:]
    aad = _make_aad(row_id, column_name)
    return aes_256_gcm_decrypt(cell_key.data, nonce, ciphertext, aad)
