import base64
from typing import Dict, Optional, Any


class LambdaResponse(object):
    def __init__(
        self,
        status: int,
        response: Optional[Any],
        headers: Optional[Dict[str, str]],
    ):
        self.response = {"statusCode": status}
        if response:
            if (("Content-Type" in headers and
                    headers["Content-Type"] is not None and
                    not headers["Content-Type"].startswith("application/json")) or (
                    "Content-Encoding" in headers and
                    headers["Content-Encoding"] is not None and
                    headers["Content-Encoding"].startswith("gzip")
            )):
                self.response["body"] = base64.b64encode(response)
                headers["x-lambda-base64-body"] = "true"
            else:
                self.response["body"] = response
        if headers:
            self.response["headers"] = {}
            for header in headers:
                self.response["headers"][header] = headers[header]
