import base64
import logging
from typing import Dict, Optional, Any

logger = logging.getLogger("konglambdamessagehandler")


class LambdaResponse(object):
    def __init__(
        self,
        status: int,
        response: Optional[Any],
        headers: Optional[Dict[str, str]],
    ):
        headers = headers or {}
        self.response: Dict[str, Any] = {"statusCode": status}
        if response:
            content_type = headers.get("Content-Type")
            content_encoding = headers.get("Content-Encoding")
            if (
                (content_type is not None and not content_type.startswith("application/json"))
                or (content_encoding is not None and content_encoding.startswith("gzip"))
            ):
                self.response["body"] = base64.b64encode(response)
                logger.info("Base64 Body size %s", len(self.response["body"]))
                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]
