# 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 import Logger
from aws_lambda_powertools.middleware_factory import lambda_handler_decorator
from aws_lambda_powertools.utilities.data_classes import APIGatewayProxyEvent
from czc import json

from czc.api.v2.exceptions import InternalServerError, ServiceError


logger = Logger()


@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_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
    """
    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)
        return {
            'statusCode': err.status_code,
            'body': err.body,
        }
    except Exception as err:
        logger.exception(err)
        return {
            'statusCode': HTTPStatus.INTERNAL_SERVER_ERROR,
            'body': InternalServerError(default_error_message).body
        }
