"""
Envelope encryption for HL API private keys.

Plaintext keys never touch disk: stored values are Fernet-encrypted with a key
derived from the env KEK (ENCRYPTION_KEY).

Blob formats:
  enc:v2:<salt_b64url>:<fernet_token>   — current. Per-blob random 16-byte
      salt; Fernet key = PBKDF2-HMAC-SHA256(KEK, salt, 600k iters). Derived
      keys are cached per salt so steady-state decryption pays the KDF once.
  enc:v1:<fernet_token>                 — legacy. Fernet key = single unsalted
      SHA-256(KEK). Read-only: decrypt still works, encrypt always writes v2.
      Upgrade rows with scripts/reencrypt_keys.py (H4 fix).
  <plaintext>                            — pre-encryption rows. Refused in prod.

Rotate KEK → re-encrypt all keys offline (scripts/reencrypt_keys.py).
"""
import base64
import hashlib
import logging
import os
from typing import Dict, Optional

from cryptography.fernet import Fernet, InvalidToken
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC

from app.config import settings

logger = logging.getLogger(__name__)

# OWASP-recommended floor for PBKDF2-HMAC-SHA256. The KEK is already a
# high-entropy secret, so the KDF mainly buys defence-in-depth against a
# weak / partially-leaked / reused KEK. ~0.2-0.4s, paid once per unique salt.
_PBKDF2_ITERATIONS = 600_000
_SALT_BYTES = 16

ENC_PREFIX_V1 = "enc:v1:"
ENC_PREFIX_V2 = "enc:v2:"
# Kept for older imports / scripts that reference the original name.
ENC_PREFIX = ENC_PREFIX_V1


def _check_kek(raw: str) -> str:
    if not raw or len(raw) < 32:
        raise RuntimeError(
            "ENCRYPTION_KEY must be set to at least 32 random chars (e.g. `openssl rand -hex 32`)"
        )
    return raw


def _derive_v1_key(raw: str) -> bytes:
    """Legacy: single unsalted SHA-256 of the KEK."""
    digest = hashlib.sha256(_check_kek(raw).encode("utf-8")).digest()
    return base64.urlsafe_b64encode(digest)


def _derive_v2_key(raw: str, salt: bytes) -> bytes:
    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA256(),
        length=32,
        salt=salt,
        iterations=_PBKDF2_ITERATIONS,
    )
    return base64.urlsafe_b64encode(kdf.derive(_check_kek(raw).encode("utf-8")))


_fernet_v1: Optional[Fernet] = None
# salt → Fernet. One entry per unique salt actually seen: encryption reuses a
# single process-lifetime salt, decryption adds one per distinct stored blob.
_fernet_v2_cache: Dict[bytes, Fernet] = {}
_V2_CACHE_MAX = 4096

# Salt for NEW encryptions in this process — generated once so the encrypt
# path pays the 600k-iteration KDF a single time per process, while blobs
# written by other processes/runs still carry their own salt.
_encrypt_salt: Optional[bytes] = None


def _cipher_v1() -> Fernet:
    global _fernet_v1
    if _fernet_v1 is None:
        _fernet_v1 = Fernet(_derive_v1_key(settings.encryption_key))
    return _fernet_v1


def _cipher_v2(salt: bytes) -> Fernet:
    cached = _fernet_v2_cache.get(salt)
    if cached is not None:
        return cached
    f = Fernet(_derive_v2_key(settings.encryption_key, salt))
    if len(_fernet_v2_cache) >= _V2_CACHE_MAX:
        _fernet_v2_cache.clear()
    _fernet_v2_cache[salt] = f
    return f


def encrypt_api_key(plaintext: str) -> str:
    global _encrypt_salt
    if _encrypt_salt is None:
        _encrypt_salt = os.urandom(_SALT_BYTES)
    salt_b64 = base64.urlsafe_b64encode(_encrypt_salt).decode("ascii")
    token = _cipher_v2(_encrypt_salt).encrypt(plaintext.encode("utf-8")).decode("utf-8")
    return f"{ENC_PREFIX_V2}{salt_b64}:{token}"


def decrypt_api_key(stored: str) -> str:
    if not stored:
        raise ValueError("Empty api key")

    if stored.startswith(ENC_PREFIX_V2):
        rest = stored[len(ENC_PREFIX_V2):]
        try:
            salt_b64, token = rest.split(":", 1)
            salt = base64.urlsafe_b64decode(salt_b64.encode("ascii"))
        except Exception as exc:
            raise RuntimeError("Malformed enc:v2 blob") from exc
        try:
            return _cipher_v2(salt).decrypt(token.encode("utf-8")).decode("utf-8")
        except InvalidToken as exc:
            raise RuntimeError("HL key decryption failed — wrong ENCRYPTION_KEY?") from exc

    if stored.startswith(ENC_PREFIX_V1):
        try:
            return _cipher_v1().decrypt(
                stored[len(ENC_PREFIX_V1):].encode("utf-8")).decode("utf-8")
        except InvalidToken as exc:
            raise RuntimeError("HL key decryption failed — wrong ENCRYPTION_KEY?") from exc

    # Legacy plaintext row (from before encryption was added). Refuse to use in prod.
    if settings.environment == "production":
        raise RuntimeError(
            "Found legacy-plaintext HL key; run migration script before production"
        )
    logger.warning("Reading LEGACY plaintext HL key — migrate ASAP")
    return stored


def is_current_format(stored: Optional[str]) -> bool:
    """True if the blob is already enc:v2 (used by scripts/reencrypt_keys.py)."""
    return bool(stored) and stored.startswith(ENC_PREFIX_V2)