"""Field-level encryption helpers for AMCS payloads."""

from __future__ import annotations

import base64
import os
from typing import Any

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

from amcs.canonical_json import canonical_bytes


def _aad_bytes(aad_dict: dict[str, Any]) -> bytes:
    return canonical_bytes(aad_dict)


def encrypt_field(
    plaintext_bytes: bytes,
    key: bytes,
    kid: str,
    aad_dict: dict[str, Any],
) -> dict[str, str]:
    """Encrypt a field using AES-GCM with context-bound AAD."""
    if len(key) not in (16, 24, 32):
        raise ValueError("AES-GCM key must be 16, 24, or 32 bytes")
    nonce = os.urandom(12)
    aesgcm = AESGCM(key)
    ciphertext = aesgcm.encrypt(nonce, plaintext_bytes, _aad_bytes(aad_dict))

    return {
        "alg": "AES-GCM",
        "kid": kid,
        "nonce_b64": base64.b64encode(nonce).decode("ascii"),
        "ciphertext_b64": base64.b64encode(ciphertext).decode("ascii"),
    }


def decrypt_field(
    enc_obj: dict[str, str],
    key: bytes,
    aad_dict: dict[str, Any],
) -> bytes:
    """Decrypt an encrypted field object using AES-GCM and matching AAD."""
    if enc_obj.get("alg") != "AES-GCM":
        raise ValueError("Unsupported encryption algorithm")
    if len(key) not in (16, 24, 32):
        raise ValueError("AES-GCM key must be 16, 24, or 32 bytes")

    nonce = base64.b64decode(enc_obj["nonce_b64"])
    ciphertext = base64.b64decode(enc_obj["ciphertext_b64"])

    aesgcm = AESGCM(key)
    return aesgcm.decrypt(nonce, ciphertext, _aad_bytes(aad_dict))
