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"):
                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]
