# 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 collections import namedtuple
from unittest.mock import Mock

import pytest
from toolz.dicttoolz import get_in

from czc import aws

from tests.sample_data.aws_lambda import function as sample_data
from tests.unittest_helpers import create_aws_client_exception, create_aws_client_throttled_exception

ROLE_ARN = 'foo-arn-role-role-ok'
EXTERNAL_ID = 'bar-id-bar-bar-bar-b'


def get_function_details(function, aws_provider):
    function_name = get_in(['Configuration', 'FunctionName'], function)
    function_arn = get_in(['Configuration', 'FunctionArn'], function)
    if not function_name:
        return None

    tags = aws_provider.client.list_tags(Resource=function_arn)

    return {**function, **tags}


def raise_error(counter):
    raise create_aws_client_exception(error_code=aws.THROTTLE_EXCEPTIONS[min(counter - 1, 2)])


def return_tags_data(Resource):
    if Resource == 'arn:aws:lambda:us-east-1:123456789012:function:function-name':
        return sample_data.FUNCTION_TAGS_DATA
    return sample_data.FUNCTION_TAGS_DATA2


def list_tags_side_effect_failure_except_one(Resource='placeholder'):
    list_tags_side_effect_failure_except_one.counter += 1
    counter = list_tags_side_effect_failure_except_one.counter
    if counter != 1:
        raise_error(counter)
    return return_tags_data(Resource)


def list_tags_side_effect_failure(Resource='placeholder'):
    list_tags_side_effect_failure.counter += 1
    raise_error(list_tags_side_effect_failure.counter)


def list_tags_side_effect(Resource='placeholder'):
    list_tags_side_effect.counter += 1
    counter = list_tags_side_effect.counter
    if counter <= 3:
        raise_error(counter)
    return return_tags_data(Resource)


@pytest.fixture(scope='function')
def context(mock_lambda_client, mock_aws_provider, mocker):
    pagination_throttle_backoff = mocker.patch('czc.aws.PAGINATION_THROTTLE_BACKOFF_TIME')
    pagination_throttle_backoff.return_value = 1
    context = namedtuple('context', ['mocker', 'mock_lambda_client', 'mock_aws_provider'])

    context.mocker = mocker
    context.mock_lambda_client = mock_lambda_client
    mock_aws_provider.service = 'lambda'
    context.mock_aws_provider = mock_aws_provider

    yield context


@pytest.mark.unit
def test_get_paginator_primary(context):
    return_value = [{'Functions': ['function_0'], 'nextToken': 'fake'}, {'Functions': ['function_1'], 'nextToken': 'fake'}]

    client_mock = Mock()
    client_mock.can_paginate.return_value = True
    mock_paginate = client_mock.get_paginator.return_value.paginate
    mock_paginate.side_effect = return_value

    aws._get_paginator('list_functions', client_mock, {})

    mock_paginate.assert_called_with()

    last_page = {'Functions': ['function'], 'nextToken': 'fake'}
    aws._get_paginator('list_functions', client_mock, {}, last_page=last_page)
    mock_paginate.assert_called_with(PaginationConfig={'StartingToken': 'fake'})

    return_value = [{'Functions': ['function_0'], 'NextMarker': 'fake'}, {'Functions': ['function_1']}]
    mock_paginate.side_effect = return_value
    last_page = {'Functions': ['function_0'], 'NextMarker': 'fake'}

    aws._get_paginator('list_functions', client_mock, {}, last_page=last_page)
    mock_paginate.assert_called_with(PaginationConfig={'StartingToken': 'fake'})


@pytest.mark.unit
def test_get_paginator_fall_back(context):
    mock_invokable_api = Mock()
    mock_invokable_api.side_effect = [{'Functions': ['function_0'], 'NextToken': 'fake'}, {'Functions': ['function_1']}]

    client_mock = Mock()
    client_mock.can_paginate.return_value = False
    client_mock.list_functions = mock_invokable_api

    last_page = None
    paginator = aws._get_paginator('list_functions', client_mock, {}, last_page)
    for idx, page in enumerate(paginator):
        assert page['Functions'][0] == f'function_{idx}'

    mock_invokable_api.side_effect = [{'Functions': ['function_0'], 'Marker': 'fake'}, {'Functions': ['function_1']}]

    paginator = aws._get_paginator('list_functions', client_mock, {}, last_page)
    for idx, page in enumerate(paginator):
        assert page['Functions'][0] == f'function_{idx}'

    # provide a starting page
    mock_invokable_api.side_effect = [{'Functions': ['function_1']}]
    last_page = {'Functions': ['function_0'], 'Marker': 'fake'}
    paginator = aws._get_paginator('list_functions', client_mock, {}, last_page)
    for page in paginator:
        assert page == {'Functions': ['function_1']}


@pytest.mark.unit
def test_get_paginated_results_throttled_fall_back_paginator(context):
    def invokable_side_effect(kwargs=None):
        invokable_side_effect.counter += 1
        if invokable_side_effect.counter > 2:
            return {'Functions': ['function_1']}

        if invokable_side_effect.counter == 1:
            return {'Functions': ['function_0'], 'NextToken': 'fake'}
        raise create_aws_client_exception(error_code=aws.THROTTLE_EXCEPTIONS[0])

    invokable_side_effect.counter = 0
    mock_invokable_api = Mock()
    mock_invokable_api.side_effect = invokable_side_effect

    client_mock = Mock()
    client_mock.can_paginate.return_value = False
    client_mock.list_functions = mock_invokable_api
    items = aws.get_paginated_results(client_mock, 'list_functions', 'Functions')
    assert len(items) == 2


@pytest.mark.unit
@pytest.mark.sns
def test_get_paginated_results_throttled_initially(context):
    def paginator_side_effect(first_pass=False):
        if first_pass:
            raise create_aws_client_exception(error_code=aws.THROTTLE_EXCEPTIONS[0])
        else:
            yield {'Functions': ['function_0'], 'nextToken': 'fake'}
            yield {'Functions': ['function_1']}

    client_mock = Mock()
    client_mock.get_paginator.return_value.paginate.side_effect = [
        paginator_side_effect(first_pass=True),
        paginator_side_effect(),
    ]
    client_mock.can_paginate.return_value = True
    items = aws.get_paginated_results(client_mock, 'list_functions', 'Functions')
    assert len(items) == 2


@pytest.mark.unit
@pytest.mark.sns
def test_get_paginated_results_throttled_primary_paginator(context):
    def paginator_side_effect(last_pass=False):
        if last_pass:
            yield {'Functions': ['function_1']}
        else:
            yield {'Functions': ['function_0'], 'nextToken': 'fake'}
            raise create_aws_client_throttled_exception()

    client_mock = Mock()
    client_mock.get_paginator.return_value.paginate.side_effect = [
        paginator_side_effect(),
        paginator_side_effect(last_pass=True),
    ]
    client_mock.can_paginate.return_value = True
    items = aws.get_paginated_results(client_mock, 'list_functions', 'Functions')
    assert len(items) == 2


@pytest.mark.unit
@pytest.mark.sns
def test_get_paginated_results_no_items(context):
    items = aws.get_paginated_results(context.mock_lambda_client, 'list_functions', 'Functions')
    assert len(items) == 0


@pytest.mark.unit
@pytest.mark.sns
def test_get_paginated_results_one_page(context, mock_lambda_role):
    _create_n_functions(context.mock_lambda_client, mock_lambda_role, 5)
    items = aws.get_paginated_results(context.mock_lambda_client, 'list_functions', 'Functions')
    assert len(items) == 5


@pytest.mark.unit
@pytest.mark.sns
def test_get_paginated_results_multiple_pages(context, mock_lambda_role):
    _create_n_functions(context.mock_lambda_client, mock_lambda_role, 98)
    items = aws.get_paginated_results(context.mock_lambda_client, 'list_functions', 'Functions')
    assert len(items) == 98


def _create_n_functions(client, role, n):
    for i in range(n):
        client.create_function(
            FunctionName=f'function_{i}', Runtime='python3.11', Role=role, Handler='foo', Code={'ZipFile': b''}
        )


@pytest.mark.unit
def test_assume_role_uses_default_region(mock_sts):
    session = aws.assume_role(ROLE_ARN, EXTERNAL_ID)
    assert session.region_name == 'us-east-1'  # makes with default region


@pytest.mark.unit
def test_assume_role_session_cached(mock_sts):
    old_session = aws.assume_role(ROLE_ARN, EXTERNAL_ID)
    new_session = aws.assume_role(ROLE_ARN, EXTERNAL_ID)
    assert old_session == new_session

    aws.assume_role_cache.clear()
    new_session = aws.assume_role(ROLE_ARN, EXTERNAL_ID)
    assert old_session != new_session


@pytest.mark.unit
def test_assume_role_reset_on_exception(mock_sts, mocker):
    mocker.patch('czc.aws._assume_role', side_effect=[1, Exception('mocked error'), 2])
    session1 = aws.assume_role(ROLE_ARN, EXTERNAL_ID)
    session2 = aws.assume_role(ROLE_ARN, EXTERNAL_ID)
    assert session1 == 1
    assert session2 == 2
