# 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 typing import Any

from aws_lambda_powertools.event_handler import APIGatewayRestResolver
from aws_lambda_powertools.logging import correlation_paths
from aws_lambda_powertools.utilities.typing import LambdaContext

from czc.api.v2.idempotency import idempotency_key_aware
from czc.logging import get_logger, get_tracer, log_context_from_event

from czpowertools.api_configuration import PublicApiConfiguration
from czpowertools.http_response import http_response
from czpowertools.middlewares.api_key_authorization_context_middleware import ApiKeyAuthorizationContextMiddleware
from czpowertools.middlewares.bypass_auth_middleware import BypassAuthMiddleware
from czpowertools.middlewares.cache_key_middleware import CacheKeyMiddleware
from czpowertools.middlewares.test_key_middleware import TestKeyMiddleware

default_logger = get_logger()
default_tracer = get_tracer()


class PublicApiGatewayRestResolver(APIGatewayRestResolver):
    """
    CloudZero Public API Gateway REST resolver.

    Extends AWS PowerTools APIGatewayRestResolver with CloudZero-specific middleware
    for handling public API requests including test key validation, caching, and
    API key authorization context.

    This resolver automatically applies the following middleware:
    - TestKeyMiddleware: Handles test key validation
    - CacheKeyMiddleware: Manages cache key generation
    - ApiKeyAuthorizationContextMiddleware: Validates API key authorization context

    The resolver also applies CloudZero logging, tracing, idempotency, and HTTP response
    decorators to all requests.
    """

    def __init__(
        self,
        configuration: PublicApiConfiguration | None = None,
    ):
        """
        Initialize the Public API Gateway REST resolver.

        Args:
            cors: CORS configuration. Defaults to WebCORSConfig() if None.
            debug: Enable debug mode for detailed error responses.
            serializer: Custom JSON serializer function.
            strip_prefixes: List of prefixes to strip from route paths.
            enable_validation: Enable request/response validation.
            response_validation_error_http_code: HTTP status code for validation errors.
        """
        configuration = configuration or PublicApiConfiguration()

        super().__init__(
            configuration.cors,
            configuration.debug,
            configuration.serializer,
            configuration.strip_prefixes,
            configuration.enable_validation,
            configuration.response_validation_error_http_code,
        )
        self.use(
            [
                BypassAuthMiddleware(paths=[configuration.swagger_path, *(configuration.additional_paths_to_bypass_auth)]),
                TestKeyMiddleware(),
                CacheKeyMiddleware(),
                ApiKeyAuthorizationContextMiddleware(),
                *(configuration.additional_middlewares),
            ]
        )

    def resolve(self, event: dict[str, Any], context: LambdaContext) -> dict[str, Any]:
        """
        Resolve the API Gateway event with CloudZero decorators and middleware.

        Applies the following decorators in order:
        - log_context_from_event: Sets up logging context from the event
        - capture_lambda_handler: AWS X-Ray tracing
        - inject_lambda_context: Injects lambda context for correlation IDs
        - idempotency_key_aware: Handles idempotency for duplicate requests
        - http_response: Formats responses and handles errors

        Args:
            event: The API Gateway event dictionary
            context: The Lambda context object

        Returns:
            dict: The formatted API Gateway response
        """
        super_class: APIGatewayRestResolver = super()  # type: ignore

        @log_context_from_event(
            default_logger,
            event_jmespaths={
                'method': 'httpMethod',
                'path': 'path',
            },
        )
        @default_tracer.capture_lambda_handler
        @default_logger.inject_lambda_context(correlation_id_path=correlation_paths.API_GATEWAY_REST)
        @idempotency_key_aware()
        @http_response(default_logger)
        def _inner_handler(event, context):
            return super_class.resolve(event, context)

        return _inner_handler(event, context)

    @property
    def cz_organization_id(self) -> str:
        """
        Get the CloudZero organization ID from the request context.

        This property retrieves the organization ID that was set by the
        ApiKeyAuthorizationContextMiddleware during request processing.

        Returns:
            str: The CloudZero organization ID for the current request
        """
        return str(self.context.get('cz_organization_id'))
