import json

from typing import Dict

from ._Message import Message
from ._DataType import DataType
from ._DefaultToken import DefaultToken


class SwaggerMessage(Message):
    def __init__(self, status_code, body=None, headers=None):
        # type: (int, str, Dict[str, str]) -> None
        self.status_code = status_code
        self.body = body
        self.headers = headers

        data_dict = {
            u"statusCode": status_code
        }

        if body is not None:
            data_dict[u"body"] = body

        if headers is not None:
            data_dict[u"headers"] = headers
        else:
            data_dict[u"headers"] = {}

        if u"Content-Type" not in data_dict[u"headers"] and body is not None:
            data_dict[u"headers"][u"Content-Type"] = u"application/json; charset=utf-8"
            self.headers = data_dict[u"headers"]

        data = json.dumps(data_dict).encode(u"utf-8")

        super(SwaggerMessage, self).__init__(DefaultToken.TK_SYS_ACK, DataType.swagger, data)
