# 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 json

import pytest

from czc import express_step_functions
from czc.constants import NAMESPACE
from czc.exceptions import AlreadyTrackingExecutionException, CannotFindLogGroupException, StepFunctionsNotFoundException
from czc.express_step_functions import ExpressStateMachine
from czt.unittest import MATCH, Context, MockEx

TEST_BASE_NAME = 'test-base-name'
TEST_STM_ARN = 'arn:aws:states:us-east-1:1234567890:express:SomeStateMachine'
TEST_EXECUTION_ARN = 'arn:aws:states:us-east-1:1234567890:express: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):
    with Context(
        mocker, express_step_functions, ['step_functions_client', 'tagging_api_client', 'time.sleep', 'cloudwatch']
    ) as context:
        yield context


@pytest.mark.unit
def test_start_stm_should_start_with_parameters_if_not_tracking_another_execution(context):
    MockEx(context.mock_step_functions_client.describe_state_machine).if_called_with(stateMachineArn=TEST_STM_ARN).return_value(
        TEST_STM_DESCRIPTION
    )
    stm = ExpressStateMachine(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'
            and args['execution_name'] == stm.name
            and args['start_timestamp'] == stm._epoch
        )

    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 = ExpressStateMachine(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
def test_start_stm_should_not_start_if_tracking_execution(context):
    MockEx(context.mock_step_functions_client.describe_state_machine).if_called_with(stateMachineArn=TEST_STM_ARN).return_value(
        TEST_STM_DESCRIPTION
    )
    stm = ExpressStateMachine(TEST_STM_ARN, TEST_BASE_NAME)
    stm.track_execution(TEST_EXECUTION_ARN, 'execution-name', 123455678)

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


@pytest.mark.unit
def test_start_stm_should_not_track_execution_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 = ExpressStateMachine(TEST_STM_ARN, TEST_BASE_NAME)
    stm.start(param1='param1-value', param2='param2-value')

    with pytest.raises(AlreadyTrackingExecutionException):
        stm.track_execution(TEST_EXECUTION_ARN, 'execution-name', 123455678)


@pytest.mark.unit
def test_start_stm_should_not_track_new_execution_if_tracking_another_execution(context):
    MockEx(context.mock_step_functions_client.describe_state_machine).if_called_with(stateMachineArn=TEST_STM_ARN).return_value(
        TEST_STM_DESCRIPTION
    )
    stm = ExpressStateMachine(TEST_STM_ARN, TEST_BASE_NAME)
    stm.track_execution(TEST_EXECUTION_ARN, 'execution-name', 123455678)

    with pytest.raises(AlreadyTrackingExecutionException):
        stm.track_execution('another-arn', 'execution-name', 123455678)


@pytest.mark.unit
@pytest.mark.parametrize(
    'test_stm',
    [
        pytest.param({}, id='Missing "loggingConfiguration"'),
        pytest.param({'loggingConfiguration': {}}, id='Missing "destinations"'),
        pytest.param({'loggingConfiguration': {'destinations': []}}, id='Empty "destinations"'),
        pytest.param(
            {'loggingConfiguration': {'destinations': [{'someOtherLogDestination': {'logGroupArn': 'not-what-we-want'}}]}},
            id='No "cloudWatchLogsLogGroup" destination',
        ),
    ],
)
def test_constructor_raises_if_cannot_find_loggroup(context, test_stm):
    MockEx(context.mock_step_functions_client.describe_state_machine).if_called_with(stateMachineArn=TEST_STM_ARN).return_value(
        test_stm
    )
    with pytest.raises(CannotFindLogGroupException):
        ExpressStateMachine(TEST_STM_ARN, TEST_BASE_NAME)


@pytest.mark.unit
@pytest.mark.parametrize('status_type', ['ExecutionSucceeded', 'ExecutionFailed'])
@pytest.mark.parametrize('iterations', [0, 2, 5])
def test_wait_for_complete_detects_completed_executions(context, iterations, status_type):
    MockEx(context.mock_step_functions_client.describe_state_machine).if_called_with(stateMachineArn=TEST_STM_ARN).return_value(
        TEST_STM_DESCRIPTION
    )
    stm = ExpressStateMachine(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_status_query = f'fields @timestamp, execution_arn, type, id | filter execution_arn="{TEST_EXECUTION_ARN}" and (type="ExecutionFailed" or type="ExecutionSucceeded") | sort @timestamp desc'

    expected_result = {'id': 100, 'type': status_type}

    def query_call(*args, **kwargs):
        return (
            [{**expected_result, 'id': str(expected_result['id'])}]
            if context.mock_cloudwatch.query.call_count == (iterations + 1)
            else []
        )

    context.mock_cloudwatch.query.side_effect = query_call

    result = stm.wait_for_complete(5, 0)
    assert result == expected_result
    assert context.mock_cloudwatch.query.call_count == iterations + 1
    assert (
        MockEx(context.mock_cloudwatch.query)
        .call(iterations + 1)
        .was_called_with(TEST_LOG_GROUP, expected_status_query, MATCH(lambda x: x < stm._epoch))
    )


@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 = ExpressStateMachine(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_cloudwatch.query.return_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 = ExpressStateMachine(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_status_query = (
        f'fields @timestamp, @message, execution_arn, id | filter execution_arn="{TEST_EXECUTION_ARN}" | sort id'
    )

    logs = [{'id': x} for x in range(5)]

    context.mock_cloudwatch.query.return_value = [{'@message': json.dumps(item)} for item in logs]

    result = stm.get_states()
    assert result == logs
    assert MockEx(context.mock_cloudwatch.query).was_called_with(
        TEST_LOG_GROUP, expected_status_query, MATCH(lambda x: x < stm._epoch), limit=5000
    )
    assert context.mock_cloudwatch.query.call_count == 1


@pytest.mark.unit
def test_get_states_uses_max_states(context):
    MockEx(context.mock_step_functions_client.describe_state_machine).if_called_with(stateMachineArn=TEST_STM_ARN).return_value(
        TEST_STM_DESCRIPTION
    )
    stm = ExpressStateMachine(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_status_query = (
        f'fields @timestamp, @message, execution_arn, id | filter execution_arn="{TEST_EXECUTION_ARN}" | sort id'
    )

    logs = [{'id': x} for x in range(5)]

    context.mock_cloudwatch.query.return_value = [{'@message': json.dumps(item)} for item in logs]

    result = stm.get_states(max_states=250)
    assert result == logs
    assert MockEx(context.mock_cloudwatch.query).was_called_with(
        TEST_LOG_GROUP, expected_status_query, MATCH(lambda x: x < stm._epoch), limit=250
    )


@pytest.mark.unit
@pytest.mark.parametrize('iterations', [0, 2, 5])
def test_get_states_gets_complete_set(context, iterations):
    MockEx(context.mock_step_functions_client.describe_state_machine).if_called_with(stateMachineArn=TEST_STM_ARN).return_value(
        TEST_STM_DESCRIPTION
    )
    query_mock = MockEx(context.mock_cloudwatch.query)

    stm = ExpressStateMachine(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')

    completion_status_query = f'fields @timestamp, execution_arn, type, id | filter execution_arn="{TEST_EXECUTION_ARN}" and (type="ExecutionFailed" or type="ExecutionSucceeded") | sort @timestamp desc'
    expected_status_query = (
        f'fields @timestamp, @message, execution_arn, id | filter execution_arn="{TEST_EXECUTION_ARN}" | sort id'
    )
    total_states = 100

    logs = [{'id': x + 1} for x in range(total_states)]

    query_results = [{'id': item['id'], '@message': json.dumps(item)} for item in logs]

    def query_call(*args, **kwargs):
        if args[1] == completion_status_query:
            return [{'id': str(total_states), 'type': 'ExecutionSucceeded'}]

        return query_results if context.mock_cloudwatch.query.call_count == (iterations + 2) else query_results[5:]

    context.mock_cloudwatch.query.side_effect = query_call
    result = stm.get_states(get_all=True, timeout=5, sleep_interval=0)

    assert result == logs
    assert query_mock.call(1).was_called_with(TEST_LOG_GROUP, completion_status_query, MATCH(lambda x: x < stm._epoch))
    for i in range(iterations + 1):
        assert query_mock.call(i + 2).was_called_with(
            TEST_LOG_GROUP, expected_status_query, MATCH(lambda x: x < stm._epoch), limit=5000
        )


@pytest.mark.unit
def test_get_states_get_all_timeout(context):
    MockEx(context.mock_step_functions_client.describe_state_machine).if_called_with(stateMachineArn=TEST_STM_ARN).return_value(
        TEST_STM_DESCRIPTION
    )
    query_mock = MockEx(context.mock_cloudwatch.query)

    stm = ExpressStateMachine(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')

    completion_status_query = f'fields @timestamp, execution_arn, type, id | filter execution_arn="{TEST_EXECUTION_ARN}" and (type="ExecutionFailed" or type="ExecutionSucceeded") | sort @timestamp desc'
    expected_status_query = (
        f'fields @timestamp, @message, execution_arn, id | filter execution_arn="{TEST_EXECUTION_ARN}" | sort id'
    )
    total_states = 100

    logs = [{'id': x + 1} for x in range(total_states)]

    query_results = [{'id': item['id'], '@message': json.dumps(item)} for item in logs]

    query_mock.if_called_with(TEST_LOG_GROUP, completion_status_query, MATCH(lambda x: x < stm._epoch)).return_value(
        [{'id': str(total_states), 'type': 'ExecutionSucceeded'}]
    )
    query_mock.if_called_with(TEST_LOG_GROUP, expected_status_query, MATCH(lambda x: x < stm._epoch), limit=5000).return_value(
        query_results[5:]
    )

    with pytest.raises(TimeoutError):
        stm.get_states(get_all=True, timeout=5, sleep_interval=0)


@pytest.mark.unit
def test_get_stm_arn_from_logical_id_tag_finds_stm(context):
    def match_tag_filter(tag_filter):
        return any(
            filter(lambda x: x['Key'] == 'aws:cloudformation:logical-id' and x['Values'] == [TEST_BASE_NAME], tag_filter)
        ) and any(filter(lambda x: x['Key'] == 'cz:namespace' and x['Values'] == [NAMESPACE], tag_filter))

    MockEx(context.mock_tagging_api_client.get_resources).if_called_with(TagFilters=MATCH(match_tag_filter)).return_value(
        {
            'ResourceTagMappingList': [
                {'ResourceARN': 'arn:aws:lambda:us-east-1:1234567890:handler:SomeStateMachine'},
                {'ResourceARN': TEST_STM_ARN},
            ]
        }
    )

    result = express_step_functions.get_stm_arn_from_logical_id_tag(TEST_BASE_NAME)

    assert result == TEST_STM_ARN


@pytest.mark.unit
def test_get_stm_arn_from_logical_id_tag_does_not_find_stm(context):
    context.mock_tagging_api_client.get_resources.return_value = {
        'ResourceTagMappingList': [{'ResourceARN': 'arn:aws:lambda:us-east-1:1234567890:handler:SomeStateMachine'}]
    }

    with pytest.raises(StepFunctionsNotFoundException):
        express_step_functions.get_stm_arn_from_logical_id_tag(TEST_BASE_NAME)
