"""Key encryption at rest for member chat keys (members.key_token).

Design (member feedback, 2026-10):
  Member keys are encrypted with AES-256-GCM using a key-encryption key
  (KEK) supplied via the PORTAL_KEK environment variable. The plaintext
  exists only in process memory at the moment a key is served (dashboard
  display, LibreChat key feed, LiteLLM management calls) — never on disk
  after startup. A stolen database file yields ciphertext, not spending
  power.

Format:
  "enc1:" + base64(nonce[12] || ciphertext || tag[16])
  Values without the prefix are treated as legacy plaintext (pre-migration
  rows) and passed through on decrypt — this keeps the migration idempotent
  and rollout safe.

Fail-closed behavior:
  - encrypt_key() raises if PORTAL_KEK is unset: nothing new is stored in
    plaintext.
  - decrypt_key() raises on an encrypted value without PORTAL_KEK (wrong or
    missing key = the key cannot be served).
  - get_db()'s migration encrypts any remaining plaintext rows at startup;
    it is skipped (with a CRITICAL log) if PORTAL_KEK is unset.
"""

import base64
import logging
import os

from cryptography.hazmat.primitives.ciphers.aead import AESGCM

log = logging.getLogger("uvicorn.error")

_PREFIX = "enc1:"
_AAD = b"member-chat-key-v1"


def _kek() -> bytes:
    raw = os.environ.get("PORTAL_KEK", "")
    if not raw:
        raise RuntimeError(
            "PORTAL_KEK is not set — refusing to encrypt/decrypt member keys"
        )
    key = base64.b64decode(raw)
    if len(key) != 32:
        raise RuntimeError("PORTAL_KEK must decode to exactly 32 bytes (base64)")
    return key


def encrypt_key(plain: str | None) -> str | None:
    """Encrypt a member key. Already-encrypted and empty values pass through."""
    if not plain or plain.startswith(_PREFIX):
        return plain
    nonce = os.urandom(12)
    ct = AESGCM(_kek()).encrypt(nonce, plain.encode(), _AAD)
    return _PREFIX + base64.b64encode(nonce + ct).decode()


def decrypt_key(stored: str | None) -> str | None:
    """Decrypt a member key. Legacy plaintext rows pass through unchanged."""
    if not stored:
        return stored
    if not stored.startswith(_PREFIX):
        return stored  # pre-migration plaintext
    blob = base64.b64decode(stored[len(_PREFIX):])
    pt = AESGCM(_kek()).decrypt(blob[:12], blob[12:], _AAD)
    return pt.decode()


def migrate_plaintext_keys(conn) -> int:
    """Encrypt any plaintext key_token rows. Idempotent; returns rows changed.

    Skipped with a CRITICAL log when PORTAL_KEK is unset — the portal keeps
    running on legacy plaintext until the operator sets the key.
    """
    if not os.environ.get("PORTAL_KEK"):
        log.critical(
            "PORTAL_KEK unset — member keys remain PLAINTEXT at rest; "
            "set PORTAL_KEK to a base64 32-byte key and restart"
        )
        return 0
    rows = conn.execute(
        "SELECT email, key_token FROM members "
        "WHERE key_token != '' AND key_token IS NOT NULL "
        "AND key_token NOT LIKE 'enc1:%'"
    ).fetchall()
    changed = 0
    for email, kt in rows:
        enc = encrypt_key(kt)
        if enc != kt:
            conn.execute(
                "UPDATE members SET key_token = ? WHERE email = ?", (enc, email)
            )
            changed += 1
    return changed