# 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 unittest.mock import ANY

import pytest

from czc import json
from czc.exceptions import AlreadyTrackingExecutionException
from czc.standard_step_functions import StandardStateMachine, Status
from czt.unittest import MATCH, Context, MockEx

TEST_BASE_NAME = 'test-base-name'
TEST_STM_ARN = 'arn:aws:states:us-east-1:1234567890:standard:SomeStateMachine'
TEST_EXECUTION_ARN = 'arn:aws:states:us-east-1:1234567890:standard:SomeStateMachine:3e0e0d25-b482-4e61-aa52-9994b2dcd51e.1648666481:cfcc3ad0-d474-4ad7-b5ee-1ee40bfabcea'
TEST_LOG_GROUP = 'SomeStateMachineLogGroup'
TEST_LOGS_ARN = f'arn:aws:logs:us-east-1:1234567890:log-group:{TEST_LOG_GROUP}'

TEST_STM_DESCRIPTION = {
    'loggingConfiguration': {
        'destinations': [
            {'someOtherLogDestination': {'logGroupArn': 'not-what-we-want'}},
            {'cloudWatchLogsLogGroup': {'logGroupArn': TEST_LOGS_ARN}},
        ]
    }
}


@pytest.fixture(scope='function')
def context(mocker):
    from czc import standard_step_functions

    test_states = [
        {'id': x + 1, 'previousEventId': x - 1, 'type': state[0], **({state[1]: {'name': state[0]}} if state[1] else {})}
        for x, state in enumerate(standard_step_functions.detail_map.items())
    ]

    with Context(
        mocker,
        standard_step_functions,
        ['step_functions_client', 'get_paginated_results', 'time.sleep'],
        extra={'standard_step_functions': standard_step_functions, 'test_states': test_states},
    ) as context:
        yield context


@pytest.mark.unit
def test_start_stm_should_start_with_parameters_if_not_tracking_another_execution(context):
    stm = StandardStateMachine(TEST_STM_ARN, TEST_BASE_NAME)
    stm.start(param1='param1-value', param2='param2-value')

    def match_args(x):
        args = json.loads(x)
        return args['param1'] == 'param1-value' and args['param2'] == 'param2-value'

    assert MockEx(context.mock_step_functions_client.start_execution).was_called_once_with(
        stateMachineArn=TEST_STM_ARN, name=MATCH(lambda x: x.startswith(TEST_BASE_NAME)), input=MATCH(match_args)
    )


@pytest.mark.unit
def test_start_stm_should_not_start_if_already_started_execution(context):
    MockEx(context.mock_step_functions_client.describe_state_machine).if_called_with(stateMachineArn=TEST_STM_ARN).return_value(
        TEST_STM_DESCRIPTION
    )
    stm = StandardStateMachine(TEST_STM_ARN, TEST_BASE_NAME)
    stm.start(param1='param1-value', param2='param2-value')

    with pytest.raises(AlreadyTrackingExecutionException):
        stm.start(paramA='paramA-value')


@pytest.mark.unit
@pytest.mark.parametrize('status_type', ['ExecutionSucceeded', 'ExecutionFailed'])
@pytest.mark.parametrize('iterations', [0, 2, 5])
@pytest.mark.parametrize('delayed_execution_visibility', [True, False])
def test_wait_for_complete_detects_completed_executions(context, iterations, status_type, delayed_execution_visibility):
    MockEx(context.mock_step_functions_client.describe_state_machine).if_called_with(stateMachineArn=TEST_STM_ARN).return_value(
        TEST_STM_DESCRIPTION
    )
    stm = StandardStateMachine(TEST_STM_ARN, TEST_BASE_NAME)
    context.mock_step_functions_client.start_execution.return_value = {'executionArn': TEST_EXECUTION_ARN}
    stm.start(param1='param1-value', param2='param2-value')

    expected_result = {'id': 100, 'previousEventId': 99, 'type': status_type, 'timestamp': 12345, 'details': {}}

    final_status = Status.succeeded.value if status_type == 'ExecutionSucceeded' else Status.failed.value
    total_iterations = iterations + 1 + (1 if delayed_execution_visibility else 0)

    def describe_execution(*args, **kwargs):
        nonlocal delayed_execution_visibility

        if delayed_execution_visibility:
            delayed_execution_visibility = False
            raise Exception('Execution is not visible yet')

        return {
            'status': Status.running.value
            if context.mock_step_functions_client.describe_execution.call_count < total_iterations
            else final_status
        }

    context.mock_step_functions_client.describe_execution.side_effect = describe_execution
    context.mock_step_functions_client.get_execution_history.return_value = {'events': [expected_result]}

    result = stm.wait_for_complete(5, 0)
    assert result == expected_result
    assert context.mock_step_functions_client.describe_execution.call_count == total_iterations
    assert (
        MockEx(context.mock_step_functions_client.describe_execution)
        .call(total_iterations)
        .was_called_with(executionArn=TEST_EXECUTION_ARN)
    )
    assert MockEx(context.mock_step_functions_client.get_execution_history).was_called_once_with(
        executionArn=TEST_EXECUTION_ARN, maxResults=1, reverseOrder=True
    )


@pytest.mark.unit
def test_wait_for_complete_timesout(context):
    MockEx(context.mock_step_functions_client.describe_state_machine).if_called_with(stateMachineArn=TEST_STM_ARN).return_value(
        TEST_STM_DESCRIPTION
    )
    stm = StandardStateMachine(TEST_STM_ARN, TEST_BASE_NAME)
    context.mock_step_functions_client.start_execution.return_value = {'executionArn': TEST_EXECUTION_ARN}
    stm.start(param1='param1-value', param2='param2-value')

    context.mock_step_functions_client.describe_execution.return_value = {'status': Status.running.value}

    with pytest.raises(TimeoutError):
        stm.wait_for_complete(5, 0)


@pytest.mark.unit
def test_get_states(context):
    MockEx(context.mock_step_functions_client.describe_state_machine).if_called_with(stateMachineArn=TEST_STM_ARN).return_value(
        TEST_STM_DESCRIPTION
    )
    stm = StandardStateMachine(TEST_STM_ARN, TEST_BASE_NAME)
    context.mock_step_functions_client.start_execution.return_value = {'executionArn': TEST_EXECUTION_ARN}
    stm.start(param1='param1-value', param2='param2-value')
    MockEx(context.mock_get_paginated_results).if_called_with(
        ANY, 'get_execution_history', 'events', executionArn=TEST_EXECUTION_ARN
    ).return_value(context.test_states)

    result = stm.get_states()
    assert len(result) == len(context.test_states)
    for state in result:
        if context.standard_step_functions.detail_map[state['type']] is None:
            assert state['details'] == {}
        else:
            assert state['details']['name'] == state['type']
