# Copyright (c) 2022 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 functools
from typing import Any, Optional

import simplejson as json
from aws_lambda_powertools.utilities.parser import BaseModel
from toolz.dicttoolz import dissoc

import czc.data.simulations as simulations
from czc.api.environment import TestKeyEnvironment
from czc.dynamodb import dynamo_safe, get_expiry_timestamp
from czc.logging import get_logger

logger = get_logger()

SET_SIMULATION_PARAMETER = '__set_simulation__'
MATCHES_ANY = '__matches_any__'


class SimulationParameters(BaseModel):
    test_key: str
    response: Any
    ttl_in_seconds: Optional[int]


def _get_target(handler):
    return f'function-{handler.__module__}.{handler.__name__}'


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

    simulation = simulations.get(test_key, _get_target(handler))

    if not simulation:
        return False, None

    for sim in simulation.get('simulated_responses', []):
        expected_args = json.loads(sim['args'])
        expected_kwargs = json.loads(sim['kwargs'])
        args_match = len(expected_args) == len(args) and all(
            expected_args[idx] == MATCHES_ANY or expected_args[idx] == arg for idx, arg in enumerate(args)
        )
        kwargs_match = len(expected_kwargs) == len(kwargs) and all(
            expected_kwargs.get(k) == MATCHES_ANY or expected_kwargs.get(k) == v for k, v in kwargs.items()
        )

        if args_match and kwargs_match:
            return True, json.loads(sim['response'])

    return False, None


def _set_simulation(handler, simulation: SimulationParameters, args, kwargs):
    target = _get_target(handler)
    simulation_data = simulations.get(simulation.test_key, target) or {'simulated_responses': []}

    simulation_data['simulated_responses'].append(
        {
            'args': json.dumps(args),
            'kwargs': json.dumps(dissoc(kwargs, SET_SIMULATION_PARAMETER)),
            'response': json.dumps(simulation.response),
        }
    )

    if simulation.ttl_in_seconds and simulation.ttl_in_seconds > simulation_data.get('ttl', 0):
        simulation_data['ttl'] = dynamo_safe(get_expiry_timestamp(simulation.ttl_in_seconds))

    simulations.put(simulation.test_key, _get_target(handler), simulation_data)
    return simulation.response


def handle_simulation():
    """
    ***************
    Function Simulation
    ***************
    This decorator allows for function simulation. When making a function 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 function call.

    This simulation is primarily useful when a function call will call some external service, such as an Azure or GCP API, which
    needs to be simulated to test situations that are different from our own Azure or GCP account. For example making Azure calls
    for an EnterpriseAgreement account.

    **** 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 function module and call the function that should be simulated with an additional parameter
    named `__set_simulation__`. The value must be an instance of the class `Simulation`.
    The `Simulation` class requires a test key and a response. It can optionally take arguments to match.
    When the `__set_simulation__` parameter is passed, instead of making the function 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 internal function call that makes a REST call to Azure. For the smoke test you want the function call to return a
    specific response instead of making the REST call to Azure. Then, from the smoke test code in feature A:

    1. Import the internal module/function call that invokes the Azure REST call.
    2. Generate a test key
    3. Call the module function with the additional parameter `__set_simulation__` with the test key
       and the desired response. Optionally, the simulation parameters can contain a set of args/kwargs that must be matched.
    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):
        @functools.wraps(handler)
        def handler_wrapper(*args, **kwargs):
            simulation: SimulationParameters = kwargs.get(SET_SIMULATION_PARAMETER)
            if simulation:
                return _set_simulation(handler, simulation, args, kwargs)

            try:
                use_simulation, response = _check_for_simulated_response(handler, args, kwargs)
                if use_simulation:
                    return response
            except Exception as ex:
                logger.error('Failed to execute a simulation - calling through to the original handler')
                logger.exception(ex)

            return handler(*args, **kwargs)

        return handler_wrapper

    return decorator
