# 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.exceptions import UnauthorizedError
from aws_lambda_powertools.event_handler.middlewares import BaseMiddlewareHandler, NextMiddleware
from toolz import get_in


class ApiKeyAuthorizationContextMiddleware(BaseMiddlewareHandler):
    """
    Middleware to validate API key authorization context for public API requests.

    This middleware checks that the request contains a valid CloudZero organization ID
    in the API Gateway authorizer context. If the organization ID is missing, it raises
    an UnauthorizedError. Otherwise, it adds the organization ID to the app context
    for use by other components.
    """

    params: list[str]

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

    def handler(
        self,
        app: ApiGatewayResolver,
        next_middleware: NextMiddleware,
    ) -> Response:
        """
        Process the API Gateway request and validate authorization context.

        Extracts the CloudZero organization ID from the API Gateway authorizer context
        and adds it to the application context for downstream use.

        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

        Raises:
            UnauthorizedError: If cz_organization_id is not present in the authorizer context
        """
        if app.current_event.path.endswith('/swagger'):
            return next_middleware(app)

        current_event = app.current_event
        authorizer = get_in(['requestContext', 'authorizer'], current_event)
        if 'cz_organization_id' not in authorizer:
            raise UnauthorizedError('You are not authorized to access this resource')

        app.append_context(cz_organization_id=authorizer['cz_organization_id'])
        return next_middleware(app)
