# 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 czc import json
from datetime import datetime, timedelta, timezone
from unittest.mock import ANY, Mock

import pytest
from freezegun import freeze_time

from czc import step_functions
from czc.json import serializable
from czc.step_functions import StateMachine
from czt.unittest import MATCH, Context, MockEx

NOT_RUNNING = {'executions': []}
TABLE = Mock()


def _running(stm_arn, base_name):
    example_stm = StateMachine(stm_arn, base_name)
    return {'executions': [{'name': example_stm.name}]}


@pytest.fixture(scope='function')
def context(mocker):
    with Context(mocker, step_functions, ['Lease', 'step_functions_client', 'time.sleep']) as context:
        yield context


@pytest.mark.unit
@pytest.mark.parametrize('stm_arn,lock', [
    ('arn:start-it', {'PK': 'foo', 'SK': 'bar'}),
    ('arn:do-the-thing', {'PK': 'lock_it', 'SK': 'some-org'}),
])
def test_start_stm_should_not_start_if_lease_unavailable(context, stm_arn, lock):
    context.mock_Lease.return_value.__enter__.return_value = None
    step_functions.start_stm(TABLE, stm_arn, lock, 'base_name')
    assert MockEx(context.mock_Lease).was_called_once_with(TABLE, lock, ANY)
    assert not context.mock_step_functions_client.list_executions.called
    assert not context.mock_step_functions_client.start_execution.called


@pytest.mark.unit
@pytest.mark.parametrize('stm_arn,base_name', [
    ('arn:start-it', 'some.execution'),
    ('arn:do-the-thing', 'run.the.thing'),
])
def test_start_stm_should_not_start_if_stm_already_running(context, stm_arn, base_name):
    MockEx(context.mock_step_functions_client.list_executions).if_called_with(
        stateMachineArn=stm_arn, statusFilter='RUNNING').return_value(_running(stm_arn, base_name))
    step_functions.start_stm(TABLE, stm_arn, {}, base_name)
    assert MockEx(context.mock_step_functions_client.list_executions).was_called_once_with(
        stateMachineArn=stm_arn, statusFilter='RUNNING')
    assert not context.mock_step_functions_client.start_execution.called


@pytest.mark.unit
@pytest.mark.parametrize('stm_arn,base_name,params', [
    ('arn:xyz', 'some.execution', {'foo': 'bar', 'baz': 42}),
    ('arn:ijk', 'run.the.thing', {'x': 1, 'y': 2, 'z': 3}),
    ('arn:ijk', 'needs.extended.encoder', {'obj': {1, 2, 3}}),
])
def test_start_stm_should_start_execution_for_provided_parameters(context, stm_arn, base_name, params):
    with freeze_time():
        context.mock_step_functions_client.list_executions.side_effect = [NOT_RUNNING, _running(stm_arn, base_name)]
        step_functions.start_stm(TABLE, stm_arn, {}, base_name, **params)
    assert MockEx(context.mock_step_functions_client.start_execution).was_called_once_with(
        stateMachineArn=stm_arn,
        name=MATCH(lambda x: x.startswith(base_name)),
        input=MATCH(lambda x: json.loads(x) == serializable(params))
    )


@pytest.mark.unit
@pytest.mark.parametrize('stm_arn,base_name', [
    ('arn:xyz', 'some.execution'),
    ('arn:ijk', 'run.the.thing'),
])
def test_start_stm_should_wait_for_running(context, stm_arn, base_name):
    with freeze_time():
        running = _running(stm_arn, base_name)
        context.mock_step_functions_client.list_executions.side_effect = [NOT_RUNNING] * 5 + [running]
        step_functions.start_stm(TABLE, stm_arn, {}, base_name)
    assert context.mock_step_functions_client.list_executions.call_count == 6


@pytest.mark.unit
def test_start_stm_should_not_page_all_of_time(context):
    stm_arn = 'arn:xyz'
    base_name = 'real.base.name'
    now = datetime.now(timezone.utc)
    context.mock_step_functions_client.list_executions.side_effect = [NOT_RUNNING, {
        'executions': [{'name': 'foo', 'startDate': now}],
        'nextToken': 'token1'
    }, {'executions': [{'name': 'bar', 'startDate': now - timedelta(minutes=5)}],
        'nextToken': 'token2'
        }, _running(stm_arn, base_name)]
    with freeze_time(now):
        step_functions.start_stm(TABLE, stm_arn, {}, base_name)
    assert context.mock_step_functions_client.list_executions.call_count == 4
    assert MockEx(context.mock_step_functions_client.list_executions).call(3).was_called_with(
        stateMachineArn=stm_arn,
        nextToken='token1'
    )
    assert MockEx(context.mock_step_functions_client.list_executions).call(4).was_called_with(
        stateMachineArn=stm_arn
    )
