import json
import base64


class ResponseFormat:
    """响应格式转换"""
    def __init__(self, body, status_code, is_base64_encoded=False, headers=None):
        """
        :param body: 内容
        :param status_code: 状态码
        :param is_base64_encoded: 是否需要base64加密
        :param headers: 请求头内容
        """
        self.status_code = status_code
        self.is_base64_encoded = is_base64_encoded
        if not headers:
            headers = dict()
        self.headers = headers
        self.body = body

    def to_json(self):
        """转换为json"""
        if self.is_base64_encoded:
            raise Exception("转换json，请把is_base64_encoded设置为False")
        self.headers.update({"Content-type": "application/json"})
        return {"statusCode": self.status_code,
                'isBase64Encoded': self.is_base64_encoded,
                'headers': self.headers,
                "body": json.dumps(self.body)}

    def to_base64(self):
        """转换为base64"""
        if not self.is_base64_encoded:
            raise Exception("转换base64，请把is_base64_encoded设置为True")
        self.headers.update({"Content-type": "text/html; charset=utf-8"})
        return {"statusCode": self.status_code,
                'isBase64Encoded': self.is_base64_encoded,
                'headers': self.headers,
                "body": base64.b64encode(str(self.body).encode(encoding="utf-8")).decode()}