"""Application-level PHI field encryption (AES-256-GCM).

Two key sources, selected per value at decrypt time and by environment at
encrypt time:

1. AWS KMS — set ``PHI_KMS_KEY_ID`` (key ARN or alias). Values are encrypted
   with the KMS Encrypt API; the returned CiphertextBlob is self-describing
   and supports automatic key rotation without re-encryption.
2. Local DEK — set ``PHI_LOCAL_ENCRYPTION_KEY`` (base64-encoded 32-byte key).
   Used by dev/test environments where no AWS credentials exist. Encryption
   uses AES-256-GCM directly via the ``cryptography`` package.

Ciphertext format::

    enc:v1:<mode>:<payload>

where mode is ``kms`` (base64 CiphertextBlob) or ``local``
(base64(nonce || ciphertext_with_tag)). Values without the ``enc:`` prefix
are treated as legacy plaintext on read only.
"""

import base64
import os
import threading

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

ENC_PREFIX = "enc:"
FORMAT_VERSION = "v1"
KMS_MODE = "kms"
LOCAL_MODE = "local"
_GCM_NONCE_SIZE = 12


class PhiCryptoError(Exception):
    """Raised when PHI encryption or decryption fails."""


def _b64_encode(data: bytes) -> str:
    return base64.urlsafe_b64encode(data).decode("ascii")


def _b64_decode(value: str) -> bytes:
    try:
        return base64.urlsafe_b64decode(value.encode("ascii"))
    except Exception as exc:
        raise PhiCryptoError(f"Invalid ciphertext encoding: {exc}") from exc


class _KmsBackend:
    """Lazy boto3 KMS backend (created only when PHI_KMS_KEY_ID is set)."""

    def __init__(self, key_id: str):
        import boto3  # imported lazily so dev/test never needs AWS SDK config

        self.key_id = key_id
        session = boto3.session.Session()
        self._client = session.client("kms")

    def encrypt(self, plaintext: str) -> bytes:
        try:
            response = self._client.encrypt(
                KeyId=self.key_id,
                Plaintext=plaintext.encode("utf-8"),
                EncryptionAlgorithm="AES_GCM",
            )
            return response["CiphertextBlob"]
        except Exception as exc:
            raise PhiCryptoError(f"KMS encrypt failed: {exc}") from exc

    def decrypt(self, blob: bytes) -> str:
        try:
            response = self._client.decrypt(CiphertextBlob=blob)
            return response["Plaintext"].decode("utf-8")
        except Exception as exc:
            raise PhiCryptoError(f"KMS decrypt failed: {exc}") from exc


class _LocalBackend:
    """AES-256-GCM with a locally-sourced data encryption key."""

    def __init__(self, raw_key: bytes):
        if len(raw_key) != 32:
            raise PhiCryptoError(
                "PHI_LOCAL_ENCRYPTION_KEY must decode to exactly 32 bytes"
            )
        self._aesgcm = AESGCM(raw_key)

    def encrypt(self, plaintext: str) -> bytes:
        nonce = os.urandom(_GCM_NONCE_SIZE)
        ciphertext = self._aesgcm.encrypt(nonce, plaintext.encode("utf-8"), None)
        return nonce + ciphertext  # GCM appends the tag to the ciphertext

    def decrypt(self, payload: bytes) -> str:
        if len(payload) <= _GCM_NONCE_SIZE:
            raise PhiCryptoError("Ciphertext too short")
        nonce, ciphertext = payload[:_GCM_NONCE_SIZE], payload[_GCM_NONCE_SIZE:]
        try:
            return self._aesgcm.decrypt(nonce, ciphertext, None).decode("utf-8")
        except InvalidTag as exc:
            raise PhiCryptoError(
                "Decryption failed (wrong key or corrupted value)"
            ) from exc


_lock = threading.Lock()
_backends = {}


def _get_backend(mode):
    """Return a backend for the requested mode, building it on first use."""
    with _lock:
        if mode in _backends:
            return _backends[mode]

        if mode == KMS_MODE:
            key_id = os.getenv("PHI_KMS_KEY_ID", "").strip()
            if not key_id:
                raise PhiCryptoError(
                    "Value was encrypted with AWS KMS but PHI_KMS_KEY_ID is not set"
                )
            backend = _KmsBackend(key_id)
        elif mode == LOCAL_MODE:
            local_key_b64 = os.getenv("PHI_LOCAL_ENCRYPTION_KEY", "").strip()
            if not local_key_b64:
                raise PhiCryptoError(
                    "Value was encrypted with a local key but "
                    "PHI_LOCAL_ENCRYPTION_KEY is not set"
                )
            backend = _LocalBackend(_b64_decode(local_key_b64))
        else:
            raise PhiCryptoError(f"Unknown encryption mode: {mode}")

        _backends[mode] = backend
        return backend


def _default_mode():
    if os.getenv("PHI_KMS_KEY_ID", "").strip():
        return KMS_MODE
    if os.getenv("PHI_LOCAL_ENCRYPTION_KEY", "").strip():
        return LOCAL_MODE
    raise PhiCryptoError(
        "PHI encryption is not configured. Set PHI_KMS_KEY_ID (production/staging) "
        "or PHI_LOCAL_ENCRYPTION_KEY (dev/test, base64-encoded 32-byte key)."
    )


def encrypt_value(plaintext: str) -> str:
    """Encrypt a plaintext string into the enc:v1 envelope format."""
    if not isinstance(plaintext, str) or not plaintext:
        raise PhiCryptoError("encrypt_value requires a non-empty string")
    mode = _default_mode()
    payload = _b64_encode(_get_backend(mode).encrypt(plaintext))
    return f"{ENC_PREFIX}{FORMAT_VERSION}:{mode}:{payload}"


def decrypt_value(ciphertext: str) -> str:
    """Decrypt an enc:v1 envelope back to plaintext.

    Values without the ``enc:`` prefix are returned unchanged so legacy
    plaintext rows remain readable during and after migration cutover.
    """
    if not isinstance(ciphertext, str) or not ciphertext.startswith(ENC_PREFIX):
        return ciphertext  # legacy plaintext passthrough
    parts = ciphertext.split(":", 3)
    if len(parts) != 4:
        raise PhiCryptoError("Malformed ciphertext envelope")
    _, version, mode, payload = parts
    if version != FORMAT_VERSION:
        raise PhiCryptoError(f"Unsupported ciphertext format version: {version}")
    backend = _get_backend(mode)
    return backend.decrypt(_b64_decode(payload))


def is_encrypted(value) -> bool:
    """True when the value is a fully-formed enc:v1 envelope."""
    if not isinstance(value, str) or not value.startswith(ENC_PREFIX):
        return False
    parts = value.split(":", 3)
    return len(parts) == 4 and parts[1] == FORMAT_VERSION