# Copyright (c) 2020 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 functools import partial
from typing import Union, Callable

import boto3
import simplejson as json
from voluptuous import ALLOW_EXTRA, Invalid, Optional, Schema

from czc.json import ExtendedEncoder

lambda_client = boto3.client('lambda')

RESPONSE = Schema({
    'data': lambda x: x,
    Optional('cursor'): str
}, required=True, extra=ALLOW_EXTRA)


class RemoteApiError(Exception):
    def __init__(self, feature, api, error_result):
        self.feature = feature
        self.api = api
        self.error = error_result

    def __str__(self):
        error_type = self.error.get('errorType', 'unknown')
        error_message = self.error.get('errorMessage', 'Unknown error.')
        return f'Remote API {self.feature}.{self.api} failed with: ({error_type}) {error_message}'


def _resolve_callable(func_or_value, *args, **kwargs):
    return func_or_value(*args, **kwargs) if callable(func_or_value) else func_or_value


def _get_function_name(namespace, feature, api, account_id):
    feature_namespace = _resolve_callable(namespace, feature)
    feature_account = _resolve_callable(account_id)
    account_prefix = f'{feature_account}:function:' if feature_account else ''
    return f'{account_prefix}cz-{feature_namespace}-{feature}-iapi-{api}'


def _response_error(response):
    return bool(response.get('FunctionError'))


def _invoke(feature, api, client, function_name, data):
    payload = json.dumps(data, cls=ExtendedEncoder, iterable_as_array=True).encode()
    response = client.invoke(FunctionName=function_name, Payload=payload)

    result = json.loads(response['Payload'].read().decode())
    if _response_error(response):
        raise RemoteApiError(feature, api, result)

    return result


def _get_paginated_result(fn, data, cursor):
    while(data):
        for x in data:
            yield x
        if cursor:
            result = fn({'cursor': cursor})
            cursor = result.get('cursor')
            data = result.get('data')
        else:
            data = []


def _call_api(fn, payload):
    result = fn(payload)
    try:
        valid = RESPONSE(result)
    except Invalid:
        return result

    cursor = valid.get('cursor')
    if cursor:
        data = valid['data']
        return _get_paginated_result(fn, data, cursor)
    else:
        return valid['data']


def invoked(namespace: Union[Callable, str], feature: str, api: str, account_id: Union[Callable, str] = None):
    """Decorator for invoking CloudZero standard internal API lambdas.

    Construct the lambda function name f'cz-{namespace}-{feature}-iapi-{api}'.

    If *account_id* is not None, then the lambda function name will be f'{account_id}:cz-{namespace}-{feature}-iapi-{api}'.

    *namespace* and *account_id* arguments can be callables if the values need to resolve at runtime.

    This decorator synchronously invokes the lambda function. The response payload must meet the minimum internal
    API schema: {'data': Any, Optional('cursor'): str}.

    If *cursor* is present in the intial response, this decorator automatically paginates by returning a generator
    of results until the server does not return a *cursor*.

    If *cursor* is _not_ present in the initial response, this decorator simply returns the *data* key from
    the response.

    The decorated function is a convenience stub for the calling client code. This decorator simply includes
    the kwargs of the decorated function in the payload to the invoked lambda function.

    Example usage:

    @invoked('alfa', 'czc', 'add-numbers')
    def stub(x=0, y=0):
        pass

    The decorator will invoke a function called 'cz-alfa-czc-iapi-add-numbers' with the payload {'x': 0, 'y': 0}.
    """
    def decorator(handler):

        def handler(**kwargs):
            function_name = _get_function_name(namespace, feature, api, account_id)
            fn = partial(_invoke, feature, api, lambda_client, function_name)
            return _call_api(fn, kwargs)

        return handler

    return decorator


def call_iapi(payload: object, namespace: Union[Callable, str], feature: str,
              api: str, account_id: Union[Callable, str] = None):
    """
    Function for invoking CloudZero standard internal API lambda functions.

    Automatically synchronously call and Construct the appropriate lambda function name according to our
    internal API conventions.
    The response payload must meet the minimum internal API schema:
        {
            'data': Any,
            Optional('cursor'): str
        }

    Construction of the iAPI function name:
        By default the function name format is:
            f'cz-{namespace}-{feature}-iapi-{api}'

        If *account_id* is provided for cross-account invokes, then the lambda function name will be
            f'{account_id}:cz-{namespace}-{feature}-iapi-{api}'

        Both *namespace* and *account_id* arguments can be callables if the values need to resolve at runtime.

    Automatic Pagination of results:
        If a *cursor* is present in the response, this function will automatically handle the pagination of the results
        and return a iterable (generator), otherwise it returns the data as received from the called function

    Example usage:

        >> call_iapi({'x':0, 'y': 0}, 'alfa', 'czc', 'add-numbers')

        The function will invoke the iAPI function called 'cz-alfa-czc-iapi-add-numbers' with
        the payload {'x': 0, 'y': 0}, returning the data contents of the response or a iterable

    Args:
        payload:
        namespace:
        feature:
        api:
        account_id:

    Returns:
        object or iterable

    """

    function_name = _get_function_name(namespace, feature, api, account_id)
    fn = partial(_invoke, feature, api, lambda_client, function_name)
    return _call_api(fn, payload)
