from functools import wraps
from fastapi import HTTPException
from schema.response_schema import APIResponse
from encryption.encryption import AESCipher


def encrypt_response(func):
    """
    返回值加密装饰器：
    - 仅对 APIResponse 的 data 字段加密（code/message 不加密，便于前端判断基础状态）
    - 若 data 为 None（如注册接口），不加密
    """

    @wraps(func)  # 保留原函数元信息（如 __name__、__doc__，避免 FastAPI 路由异常）
    async def wrapper(*args, **kwargs):
        try:
            # 1. 执行原接口函数，获取返回值（APIResponse 对象）
            response: APIResponse = await func(*args, **kwargs)

            # 2. 仅当 data 不为 None 时加密
            if response.data is not None:
                # 加密 data 字段（字典类型）
                encrypted_result = AESCipher.encrypt(response.data)
                # 替换原 data 为加密后的数据（包含密文和 IV）
                response.data = {
                    "is_encrypted": True,  # 标记是否加密，便于前端处理
                    **encrypted_result
                }

            return response

        except Exception as e:
            # 加密过程异常时，返回 500 错误
            raise HTTPException(
                status_code=500,
                detail=f"返回值加密失败：{str(e)}"
            ) from e

    return wrapper