# Copyright (c) 2017 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 _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)
    return json.loads(response['Payload'].read().decode())


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
