# Copyright (c) 2022 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
import functools
from typing import Optional

from aws_lambda_powertools import Tracer
from aws_lambda_powertools.utilities.data_classes import APIGatewayProxyEvent
from voluptuous import REMOVE_EXTRA, Invalid, Schema

from czc import json
from czc.api.adapters.response import allow_origin_response, http_response
from czc.api.v2.exceptions import BadRequestError
from czc.api.v2.testing import web_test_key_aware
from czc.authorization import organization_path_authorizer
from czc.logging import Logger, get_logger, get_tracer, log_context_from_event

default_logger = get_logger()
default_tracer = get_tracer()


def web_api(
    default_error_message: str = None,
    path_schema: dict = None,
    query_schema: Schema = None,
    multi_value_query_schema: Schema = None,
    body_schema: Schema = None,
    tracer: Tracer = default_tracer,
    logger: Logger = default_logger,
    logger_event_jmespaths: Optional[dict[str, str]] = None,
    logger_response_jmespaths: Optional[dict[str, str]] = None,
    wrap_response: bool = True,
):
    """
    Make this Lambda Handler Compliant with our Web APIs by composing all necessary decorators

    Arguments:

        default_error_message: str, optional - Default error message in case of an exception
        path_schema: dict, optional - Dictionary of path parameters to validate in addition to `organizationId`
        query_schema: Schema, optional - Query string parameters schema to validate
        multi_value_query_schema: Schema, optional - Multi-value query string parameters schema to validate
        body_schema: Schema, optional - Body schema to validate
        tracer: Tracer, optional - Tracer to use for capturing spans
        logger: Logger, optional - Logger to use for logging
        logger_event_jmespaths: dict[str, str], optional - Dictionary of event keys to extract for logging context
        wrap_response: bool, optional - Whether to wrap the response in `{'body': response}` or not.
            If the lambda handler needs to return both a body and headers for use in the `http_response` decorator, set this to `False`.

    Returns:
            Callable - Decorated Lambda Handler

    """

    def _internal_api_handler(handler):
        @log_context_from_event(
            logger=logger, event_jmespaths=logger_event_jmespaths, response_jmespaths=logger_response_jmespaths
        )
        @tracer.capture_lambda_handler
        @web_test_key_aware()
        @allow_origin_response(r'^https://.*\.cloudzero\.com$')
        @http_response(default_error_message)
        @organization_path_authorizer
        @functools.wraps(handler)
        def handler_wrapper(event, context, **kwargs):
            valid_path_schema = Schema({'organizationId': str}, required=True).extend(path_schema or {})
            valid_query_schema = query_schema or Schema({}, extra=REMOVE_EXTRA)
            valid_mvq_schema = multi_value_query_schema or Schema({}, extra=REMOVE_EXTRA)
            gateway_event = APIGatewayProxyEvent(event)
            try:
                path_params = valid_path_schema(gateway_event.path_parameters or {})
                query_params = valid_query_schema(gateway_event.query_string_parameters or {})
                mvq_params = valid_mvq_schema(gateway_event.multi_value_query_string_parameters or {})
                body_param = dict(body=body_schema(json.loads(gateway_event.body))) if body_schema else {}
            except Invalid as e:
                logger.exception(e)
                raise BadRequestError('Failed to validate parameters', invalid_exception=e)
            result = handler(event, context, **path_params, **query_params, **mvq_params, **body_param, **kwargs)
            return {'body': result} if wrap_response else result

        return handler_wrapper

    return _internal_api_handler
