import json
from datetime import datetime
from functools import wraps
from typing import Any

from encryption.encryption import aes_encrypt
from schema.response_schema import APIResponse
from pydantic import BaseModel


def encrypt_response(field: str = "data"):
    """接口返回值加密装饰器：正确序列化自定义对象为JSON"""

    def decorator(func):
        @wraps(func)
        async def wrapper(*args, **kwargs):
            original_response: APIResponse = await func(*args, **kwargs)
            field_value = getattr(original_response, field)

            if not field_value:
                return original_response

            # 自定义JSON序列化函数：处理Pydantic模型和datetime
            def json_default(obj: Any) -> Any:
                if isinstance(obj, BaseModel):
                    return obj.model_dump() 
                if isinstance(obj, datetime):
                    return obj.isoformat()
                return str(obj) 

            # 使用自定义序列化函数、确保生成标准JSON
            field_value_json = json.dumps(field_value, default=json_default)
            encrypted_data = aes_encrypt(field_value_json)
            setattr(original_response, field, encrypted_data)

            return original_response

        return wrapper

    return decorator