import json
from functools import wraps
from schema.response_schema import APIResponse
from utils.encrypt_utils import aes_encrypt


def encrypt_response(field: str = "data"):
    """接口返回值加密装饰器：默认加密APIResponse的data字段"""

    def decorator(func):
        @wraps(func)  # 保留原函数元信息（避免FastAPI路由异常）
        async def wrapper(*args, **kwargs):
            # 执行原接口函数，获取原始响应（APIResponse对象）
            original_response: APIResponse = await func(*args, **kwargs)

            # 若需加密的字段为空，直接返回原响应（如注册接口data=None）
            if not getattr(original_response, field):
                return original_response

            # 复杂数据转JSON字符串（支持datetime、字典、列表等）
            field_value = getattr(original_response, field)
            field_value_json = json.dumps(field_value, default=str)  # 处理特殊类型

            # AES加密并替换原字段
            encrypted_data = aes_encrypt(field_value_json)
            setattr(original_response, field, encrypted_data)

            # 返回加密后的响应
            return original_response

        return wrapper

    return decorator