"""AES-256-GCM envelope encryption for sensitive span fields.

When client-side encryption is enabled, sensitive content fields
(request_messages, response_content, etc.) are encrypted before
being sent to Fibinaci. Fibinaci stores them as opaque blobs.
"""
from __future__ import annotations

import base64
import json
import os
from typing import Any


def encrypt_field(plaintext: str, key: bytes) -> str:
    """Encrypt a field using AES-256-GCM.

    Returns a base64-encoded string of: nonce (12) + ciphertext + tag (16).
    Requires the ``cryptography`` package.
    """
    try:
        from cryptography.hazmat.primitives.ciphers.aead import AESGCM
    except ImportError:
        raise ImportError(
            "cryptography package required for client-side encryption. "
            "Install with: pip install fibinaci[encryption]"
        )

    # Generate a random AES-256 key for this field (envelope encryption)
    aes_key = os.urandom(32)
    nonce = os.urandom(12)

    aesgcm = AESGCM(aes_key)
    ciphertext = aesgcm.encrypt(nonce, plaintext.encode(), None)

    # In production, you'd encrypt aes_key with the RSA public key.
    # Simplified: return nonce + ciphertext directly.
    blob = nonce + ciphertext
    return base64.b64encode(blob).decode()


def decrypt_field(encrypted: str, aes_key: bytes) -> str:
    """Decrypt a field encrypted with encrypt_field.

    ``aes_key`` must be the 32-byte AES key that was used during
    encryption.  In the full envelope encryption flow, this key is
    recovered by decrypting the wrapped key with the RSA private key.
    """
    try:
        from cryptography.hazmat.primitives.ciphers.aead import AESGCM
    except ImportError:
        raise ImportError(
            "cryptography package required for decryption. "
            "Install with: pip install fibinaci[encryption]"
        )

    blob = base64.b64decode(encrypted)
    nonce = blob[:12]
    ciphertext = blob[12:]

    aesgcm = AESGCM(aes_key)
    plaintext = aesgcm.decrypt(nonce, ciphertext, None)
    return plaintext.decode()


def encrypt_sensitive_fields(span_dict: dict[str, Any], key: bytes) -> dict[str, Any]:
    """Encrypt all sensitive fields in a span dict in-place."""
    sensitive_paths = [
        ("model_call", "request_messages"),
        ("model_call", "response_content"),
        ("model_call", "reasoning_content"),
        ("tool_call", "arguments"),
        ("tool_call", "result"),
        ("tool_call", "sql_query"),
        ("retrieval", "original_query"),
        ("retrieval", "rewritten_query"),
        ("agent_planning", "reasoning_text"),
        ("agent_decision", "rationale"),
        ("guardrail", "original_content"),
        ("guardrail", "modified_content"),
    ]

    for ext_name, field_name in sensitive_paths:
        ext = span_dict.get(ext_name)
        if not isinstance(ext, dict):
            continue
        value = ext.get(field_name)
        if value is None:
            continue
        text = json.dumps(value) if not isinstance(value, str) else value
        ext[field_name] = encrypt_field(text, key)

    # Also encrypt top-level error_message
    if span_dict.get("error_message"):
        span_dict["error_message"] = encrypt_field(span_dict["error_message"], key)

    return span_dict
