# Copyright (c) 2021 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 re

from aws_lambda_powertools.event_handler.exceptions import ServiceError as PowertoolsServiceError

from czc import json
from czc.api.v2.exceptions import ForbiddenError, ServiceError
from czc.logging import get_logger
from czc.util import deep_get

logger = get_logger()


def http_response(default_error_message=None):
    """Make this Lambda Handler handle legacy response dictionaries and

    Parameters
    ----------

    Returns
    -------
    Callable
        Decorated Lambda Handler
    """

    def http_response_handler(handler):
        def handler_wrapper(event, context, **kwargs):
            try:
                res = handler(event, context, **kwargs)
                if not isinstance(res, dict):
                    raise Exception(f'Unsupported return type {type(res)}; response must be dict.')
                return {
                    'headers': res.get('headers', {}),
                    'statusCode': res.get('statusCode', 200),
                    'body': json.dumps(res['body'], iterable_as_array=True) if 'body' in res else None,
                }
            except PowertoolsServiceError as err:
                if err.status_code >= 500:
                    logger.exception(err)
                else:
                    logger.warning(err)

                return {
                    'statusCode': err.status_code,
                    'body': json.dumps(err.body, iterable_as_array=True),
                }
            except ServiceError as err:
                if err.status_code >= 500:
                    logger.exception(err)
                else:
                    logger.warning(err)
                return {
                    'statusCode': err.status_code,
                    'body': err.body,
                }
            except Exception as err:
                logger.exception(err)
                # Try and handle HTTPResponseException like objects
                if hasattr(err, 'statusCode') and hasattr(err, 'body'):
                    return {
                        'statusCode': err.statusCode,
                        'body': json.dumps(err.body, iterable_as_array=True),
                    }
                return {
                    'statusCode': 500,
                    'body': json.dumps({'error': {'message': default_error_message or 'API Error'}}, iterable_as_array=True),
                }

        return handler_wrapper

    return http_response_handler


def allow_origin_response(*origins):
    """Decorator that will check that the event.headers.origin is in origins; if the origin
    is valid, this decorator will add it to the response headers.

    Args:
        handler (func): a handler function with the signature (event, context) -> result

    Returns:
        handler (func): a lambda handler function that is authorized
    """

    def allow_origin_handler(handler):
        def handler_wrapper(event, context, **kwargs):
            logger.debug(f'Checking origin for event: {event}')

            # Check Origin
            request_origin = deep_get(event, 'headers', 'origin', ignore_case=True)
            if not any(re.match(o, str(request_origin)) for o in origins):
                logger.warning(f'Invalid request origin: {request_origin}')
                raise ForbiddenError('Unknown origin.')

            # call handler
            kwargs['request_origin'] = request_origin
            response = handler(event, context, **kwargs)

            if not isinstance(response, dict):
                raise Exception(f'Unsupported response type {type(response)}; response must be dict for *_response decorators.')

            # add origin to response headers
            current_headers = response.get('headers', {})
            cors_headers = {'Access-Control-Allow-Origin': request_origin, 'Access-Control-Allow-Credentials': 'true'}
            response['headers'] = {**current_headers, **cors_headers}
            return response

        return handler_wrapper

    return allow_origin_handler
