# Copyright (c) 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 aws_lambda_powertools.event_handler import ApiGatewayResolver, Response
from aws_lambda_powertools.event_handler.middlewares import BaseMiddlewareHandler, NextMiddleware

from czc import json


class CacheKeyMiddleware(BaseMiddlewareHandler):
    """
    Middleware to add caching headers to API responses.

    This middleware automatically generates an ETag header based on the SHA256 hash
    of the response body and applies cache control headers to enable proper HTTP caching.
    """

    params: list[str]

    def __init__(self) -> None:
        """Initialize the Cache Key middleware."""
        super().__init__()

    def handler(
        self,
        app: ApiGatewayResolver,
        next_middleware: NextMiddleware,
    ) -> Response:
        """
        Process the response and add caching headers.

        Generates an ETag based on the SHA256 hash of the response body and adds
        appropriate Cache-Control headers to the response.

        NOTE:This handler is called by the AWS Lambda Powertools framework.

        Args:
            app: The API Gateway resolver instance
            next_middleware: The next middleware in the chain

        Returns:
            Response: The response with added caching headers
        """
        response: Response = next_middleware(app)
        response_str = json.dumps(response.body)
        etag = sha256(response_str.encode()).hexdigest()
        updated_headers = {
            'Cache-Control': app.current_event.headers.get('cache-control', 'no-store'),
            'ETag': f'"{etag}"',
            **response.headers,
        }
        response.headers = updated_headers
        return response
