# 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
import collections.abc
import functools
import typing
from collections.abc import Callable
from functools import partial
from typing import Any, TypeVar, Union, cast

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

import czc.data.simulations as simulations
import czc.function_simulation as function_simulation
from czc.api.environment import TestKeyEnvironment
from czc.json import ExtendedEncoder
from czc.logging import get_logger

logger = get_logger()
lambda_client = boto3.client('lambda')


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

F = TypeVar('F', bound=Callable[..., Any])


@deprecated(reason='Use `czc.function_simulation.SimulationParameters` as it enables parameter matching.')
class Simulation(function_simulation.SimulationParameters):
    pass


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 _invoke_event(feature, api, client, function_name, data):
    payload = json.dumps(TestKeyEnvironment.add_test_key_parameter(data), cls=ExtendedEncoder, iterable_as_array=True).encode()

    response = client.invoke(FunctionName=function_name, InvocationType='Event', Payload=payload)

    if response['StatusCode'] >= 300:
        raise RemoteApiError(feature, api, response)

    return response['StatusCode']


def _coerce_to_type(type, val):  # noqa: A002
    def is_instance():
        try:
            return isinstance(val, type)
        except Exception:
            return False

    def try_coerce():
        try:
            return Coerce(type)(val)
        except Exception:
            return val

    return val if type is None or is_instance() else try_coerce()


def _is_iterable(type):  # noqa: A002
    return type is not None and typing.get_origin(type) is collections.abc.Iterable


def _get_paginated_result(fn, data, return_type, cursor, test_key):
    def get_type_arg():
        if return_type is None:
            return None

        args = typing.get_args(return_type)
        if not (_is_iterable(return_type) and len(args) == 1):
            raise TypeError('Paginated IAPI must return Iterable[<T>]')
        return args[0]

    element_return_type = get_type_arg()
    while data:
        for x in data:
            yield _coerce_to_type(element_return_type, x)
        if cursor:
            result = fn(TestKeyEnvironment.add_test_key_parameter({'cursor': cursor}, test_key))
            cursor = result.get('cursor')
            data = result.get('data')
        else:
            data = []


def _call_api(fn, payload, return_type=None):
    result = fn(TestKeyEnvironment.add_test_key_parameter(payload))

    try:
        valid = RESPONSE(result)
    except Invalid:
        return _coerce_to_type(return_type, result)

    cursor = valid.get('cursor')
    if cursor or _is_iterable(return_type):
        data = valid['data']
        return _get_paginated_result(fn, data, return_type, cursor, TestKeyEnvironment.get())
    return _coerce_to_type(return_type, valid['data'])


def _check_for_simulated_response(api):
    test_key = TestKeyEnvironment.get()
    if not test_key:
        return False, None

    simulation = simulations.get(test_key, f'iapi-{api}')
    use_simulation = simulation is not None
    return use_simulation, simulation['response'] if use_simulation else None


def _set_simulation(namespace, feature, api, account_id, simulation: Simulation):
    simulations.put(simulation.test_key, f'iapi-{api}', simulation.response)
    return simulation.response


def handle_simulation(namespace, feature, api, account_id):
    """
    ***************
    IAPI Simulation
    ***************
    This decorator allows for IAPI simulation. When making an IAPI call from a lambda and a test key is present,
    a check is made to the feature's DynamoDb table for a matching simulation record. If a matching record is found,
    the specified response will be returned instead of making the IAPI call.

    **** Setting up Simulation ***
    To set up the simulation, the smoke tests must add the simulation records to the feature's DynamoDb FROM THE
    SMOKE TEST CODE itself. This can be done using the same client functions wrapped with this decorator.

    In the smoke test code itself, import the IAPI client for whatever service the lambda under test will call.
    Using this client from the smoke test call the function that should be simulated with an additional parameter
    named `__set_simulation__` can be passed. The value must be an instance of the class `Simulation`.
    The `Simulation` class requires a test key and a response. When this additional parameter is passed, instead of making
    the IAPI call, it will store the simulation data in the database.

    For example: Suppose you are smoke testing feature A and the particular lambda feature A you wish to test
    makes an IAPI call to feature B, but for the smoke test you want the IAPI call to not get invoked but return a
    specific response. Then, from the smoke test code in feature A:

    1. Import the IAPI client for feature B.
    2. Generate a test key
    3. Call the IAPI function in the client for feature B with the additional parameter `__set_simulation__` with the test key
       and the desired response.
    4. Call the lambda in feature A that is under test passing the same smoke test key

    When the lambda in feature A calls the same IAPI for feature B, it will find the simulation and immediately return the
    specified response.

    Example:
    # Given
    @invoked(namespace, 'some-feature', 'function-name')
    def some_iapi_call(param1, param2):
        pass

    # Setup simulation
    test_key = f'test-key-{uuid4()}'
    some_iapi_call(param1='x', param2='y', __set_simulation__=Simulation(test_key=test_key, response='z'))

    ... run smoke test


    ***********
    """

    def decorator(handler):
        def handler_wrapper(**kwargs):
            @function_simulation.handle_simulation()
            def _internal_wrapper(**kwargs):
                return handler(**kwargs)

            simulation = kwargs.get(function_simulation.SET_SIMULATION_PARAMETER)
            if simulation and isinstance(simulation, Simulation):
                kwargs = {
                    **{k: function_simulation.MATCHES_ANY for k in kwargs if k != function_simulation.SET_SIMULATION_PARAMETER},
                    function_simulation.SET_SIMULATION_PARAMETER: simulation,
                }

            return _internal_wrapper(**kwargs)

        return handler_wrapper

    return decorator


def invoked(
    namespace: Union[Callable, str],
    feature: str,
    api: str,
    account_id: Union[Callable, str] = None,
    client: object = None,
    use_return_type=True,
) -> Callable[[F], F]:
    """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}.

    ***************
    IAPI Simulation
    ***************
    This decorator includes the `@handle_simulation` to enable IAPI simulation. See the documentation for
    `handle_simulation` for more information on how simulation works.

    """
    client = client or lambda_client

    def decorator(proto: F) -> F:
        def get_return_type():
            hints = typing.get_type_hints(proto)
            return hints.get('return')

        return_type = get_return_type() if use_return_type else None

        @functools.wraps(proto)
        @handle_simulation(namespace, feature, api, account_id)
        def handler(**kwargs):
            function_name = _get_function_name(namespace, feature, api, account_id)
            fn = partial(_invoke, feature, api, client, function_name)
            return _call_api(fn, kwargs, return_type)

        return cast(F, handler)

    return decorator


def event_invoked(
    namespace: Union[Callable, str], feature: str, api: str, account_id: Union[Callable, str] = None, client: object = None
):
    """Decorator for invoking CloudZero internal API lambda without waiting for a response.

    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 invokes the lambda function with InvocationType=Event, i.e. asynchronously. It does not return
    a response as it expects the client to want to "fire and forget".

    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:

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

    The decorator will send an event to a function called 'cz-alfa-czc-iapi-add-numbers' with the payload {'x': 0, 'y': 0}.

    ***************
    IAPI Simulation
    ***************
    This decorator includes the `@handle_simulation` to enable IAPI simulation. See the documentation for
    `handle_simulation` for more information on how simulation works.

    """
    client = client or lambda_client

    def decorator(_):
        @handle_simulation(namespace, feature, api, account_id)
        def handler(**kwargs):
            function_name = _get_function_name(namespace, feature, api, account_id)
            return _invoke_event(feature, api, client, function_name, kwargs)

        return handler

    return decorator


def call_iapi(
    payload: object,
    namespace: Union[Callable, str],
    feature: str,
    api: str,
    account_id: Union[Callable, str] = None,
    client: object = 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:
        client:

    Returns:
        object or iterable

    """
    client = client or lambda_client

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