# Copyright (c) 2021 CloudZero - ALL RIGHTS RESERVED - PROPRIETARY AND CONFIDENTIAL
# Unauthorized copying of this file and/or project, via any medium is strictly prohibited.
# Direct all questions to legal@cloudzero.com

from hashlib import sha256
from typing import Callable

from aws_lambda_powertools.middleware_factory import lambda_handler_decorator
from aws_lambda_powertools.utilities.data_classes import APIGatewayProxyEvent

from czc import json
from czc.api.v2.gateway import Response


@lambda_handler_decorator(trace_execution=True)
def cache_header_aware(handler: Callable[..., Response], event: APIGatewayProxyEvent, context, **kwargs):
    """Make this Lambda Handler return Cache-Control and ETag Headers

    Parameters
    ----------

    Returns
    -------
    Callable
        Decorated Lambda Handler
    """
    response = handler(event, context, **kwargs)
    response_str = json.dumps(response.body)
    etag = sha256(response_str.encode()).hexdigest()
    updated_headers = {
        'Cache-Control': event.get_header_value('cache-control', 'no-store'),
        'ETag': f'"{etag}"',
        **response.headers
    }

    return Response(
        response.body,
        response.status_code,
        updated_headers,
    )
