from aws_lambda_powertools.event_handler.api_gateway import (
    APIGatewayHttpResolver,
    Response,
)
from aws_lambda_powertools.event_handler.middlewares import (
    BaseMiddlewareHandler,
    NextMiddleware,
)


class CorrelationIdMiddleware(BaseMiddlewareHandler):
    def __init__(self, header: str):
        super().__init__()
        self.header = header

    def handler(
        self, app: APIGatewayHttpResolver, next_middleware: NextMiddleware
    ) -> Response:
        # BEFORE logic
        request_id = app.current_event.request_context.request_id
        correlation_id = app.current_event.headers.get(self.header, request_id)

        # Call next middleware or route handler ('/todos')
        response = next_middleware(app)

        # AFTER logic
        response.headers[self.header] = correlation_id

        return response