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

from czc.api.environment import TestKeyEnvironment
from czc.api.v2.constants import CLOUDZERO_TEST_KEY_HEADER


class TestKeyMiddleware(BaseMiddlewareHandler):
    """
    Middleware to handle CloudZero test key environment context.

    This middleware extracts the test key from request headers and applies
    the appropriate test environment context using CloudZero's TestKeyEnvironment.
    This allows requests to be processed in a test environment when a test key is provided.
    """

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

    def handler(
        self,
        app: ApiGatewayResolver,
        next_middleware: NextMiddleware,
    ) -> Response:
        """
        Process the request within a test key environment context.

        Extracts the test key from the request headers and applies the corresponding
        test environment context for the duration of the request processing.

        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 from the next middleware processed within the test environment
        """
        event = app.current_event

        with TestKeyEnvironment(event.headers.get(CLOUDZERO_TEST_KEY_HEADER)):
            return next_middleware(app)
