# 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

from dataclasses import dataclass, field
from http import HTTPStatus
from typing import Any, Callable, Dict, Union

from aws_lambda_powertools.middleware_factory import lambda_handler_decorator
from aws_lambda_powertools.utilities.data_classes import APIGatewayProxyEvent
from toolz.dicttoolz import get_in
from voluptuous import REMOVE_EXTRA, Invalid, Schema
from voluptuous.schema_builder import ALLOW_EXTRA

from czc import domain_events as de
from czc import json
from czc.api.v2.exceptions import (BadRequestError, InternalServerError,
                                   ServiceError, UnauthorizedError)
from czc.logging import get_logger, get_tracer

logger = get_logger()
tracer = get_tracer()


@dataclass(frozen=True)
class Response:
    # Required
    body: Dict[str, Any]
    # Optional
    status_code: int = HTTPStatus.OK
    headers: Dict[str, Any] = field(default_factory=dict)

    def __post_init__(self):
        if not isinstance(self.body, dict):
            raise ValueError('body must be dictionary ')
        if not all(isinstance(k, str) for k in self.body):
            raise ValueError('all root keys of body must be strings')


@lambda_handler_decorator(trace_execution=True)
def coerce_event(handler: Callable[..., Union[dict, Response]], event: Dict, context, **kwargs):
    """Make this Lambda Handler accept an APIGatewayProxyEvent

    Parameters
    ----------

    Returns
    -------
    Callable
        Decorated Lambda Handler
    """
    return handler(APIGatewayProxyEvent(event), context, **kwargs)


@lambda_handler_decorator(trace_execution=True)
def coerce_response(handler: Callable[..., Union[dict, Response]], event: APIGatewayProxyEvent, context, *, default_error_message: str = 'API Error', **kwargs):
    """Make this Lambda Handler return a Response Object

    Parameters
    ----------
    default_error_message: str, optional
        message for error response, by default 'API Error'

    Returns
    -------
    Callable
        Decorated Lambda Handler
    """
    raw_response = handler(event, context, **kwargs)
    if isinstance(raw_response, Response):
        return raw_response
    if isinstance(raw_response, dict):
        return Response(raw_response)
    raise Exception(f'Unsupported return type {type(raw_response)}; response body must be Union[dict, Response]')


@lambda_handler_decorator(trace_execution=True)
def serialize_response(handler: Callable[..., Response], event: APIGatewayProxyEvent, context, *, default_error_message: str = 'API Error', **kwargs):
    """Make this Lambda Handler return API Gateway and CloudZero API V2 Compliant
    responses and errors.

    Parameters
    ----------
    default_error_message: str, optional
        message for error response, by default 'API Error'

    Returns
    -------
    Callable
        Decorated Lambda Handler
    """
    exception_response = None
    try:
        response = handler(event, context, **kwargs)
        return {
            'headers': response.headers,
            'statusCode': response.status_code,
            'body': json.dumps(response.body),
        }
    except ServiceError as err:
        logger.exception(err)
        exception_response = {
            'statusCode': err.status_code,
            'body': err.body,
        }
    except Exception as err:
        logger.exception(err)
        exception_response = {
            'statusCode': HTTPStatus.INTERNAL_SERVER_ERROR,
            'body': InternalServerError(default_error_message).body
        }

    try:
        logger.info(f'Publishing exception {exception_response}')
        de.publish_event_response_data(
            default_error_message,
            'api-exception-handled',
            event.raw_event if isinstance(event, APIGatewayProxyEvent) else event,
            exception_response,
            response_filter_fn=lambda x: x
        )
    except Exception:
        logger.exception('Failed to publish api-exception-handled failure event', exc_info=True)
    return exception_response


def parameters(event_schema: Schema = None,
               path_schema: Schema = None,
               query_schema: Schema = None,
               body_schema: Schema = None,
               multi_value_query_schema: Schema = None,):
    """Add parameters to kwargs, applying a Schema if desired.

    Parameters
    ----------
    event_schema: Schema, optional
        event parameters schema, by default Schema({}, extra=REMOVE_EXTRA)
    path_schema: Schema, optional
        path parameters schema, by default Schema({}, extra=ALLOW_EXTRA)
    query_schema: Schema, optional
        query string parameters schema, by default Schema({}, extra=REMOVE_EXTRA)
    body_schema: Schema, optional
        body parameters schema, by default Schema({}, extra=REMOVE_EXTRA)
    multi_value_query_schema: Schema, optional
        multi value query string parameters schema, by default Schema({}, extra=REMOVE_EXTRA)
        NOTE: all of these values will be lists

    Returns
    -------
    Callable
        Decorated Lambda Handler
    """
    valid_event_schema = event_schema or Schema({}, extra=REMOVE_EXTRA)
    valid_path_schema = path_schema or Schema({}, extra=ALLOW_EXTRA)
    valid_query_schema = query_schema or Schema({}, extra=REMOVE_EXTRA)
    valid_body_schema = body_schema or Schema({}, extra=REMOVE_EXTRA)
    valid_multi_value_query_schema = multi_value_query_schema or Schema({}, extra=REMOVE_EXTRA)

    def decorator(handler):
        @tracer.capture_lambda_handler
        def lambda_handler(event: APIGatewayProxyEvent, context, **kwargs):
            try:
                event_params = valid_event_schema(event.raw_event or {})
                path_params = valid_path_schema(event.path_parameters or {})
                query_params = valid_query_schema(event.query_string_parameters or {})
                body_params = dict(body=valid_body_schema(event.json_body)) if body_schema else {}
                multi_value_query_params = valid_multi_value_query_schema(event.multi_value_query_string_parameters or {})
            except Invalid as e:
                raise BadRequestError('Failed to validate parameters', invalid_exception=e)
            all_params = {**event_params, **path_params, **query_params, **body_params, **multi_value_query_params}
            return handler(event, context, **all_params, **kwargs)
        return lambda_handler
    return decorator


def authorization_context(schema: Schema = None):
    """Add authorization context from event['requestContext']['authorizer'] to kwargs.

    Parameters
    ----------
    schema: schema, optional
        Voluptuous Schema, by default Schema({'cz_organization_id': str}, required=True, extra=REMOVE_EXTRA)

    Returns
    -------
    Callable
        Decorated Lambda Handler
    """
    valid_schema = schema or Schema({'cz_organization_id': str}, required=True, extra=REMOVE_EXTRA)

    def decorator(handler):
        @tracer.capture_lambda_handler
        def lambda_handler(event, context, **kwargs):
            try:
                params = valid_schema(get_in(['requestContext', 'authorizer'], event) or {})
            except Invalid:
                logger.error('Failed to access required auth information', exc_info=True)
                raise UnauthorizedError('Missing required authorization information.')
            return handler(event, context, **params, **kwargs)
        return lambda_handler
    return decorator
