# Copyright (c) 2021 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 pytest
from boto3.dynamodb.conditions import Attr, Key
from datetime import datetime, timedelta
from decimal import Decimal
from enum import Enum
from freezegun import freeze_time
from toolz.dicttoolz import update_in
from unittest.mock import Mock

from czt.unittest import Context

from czc.constants import DYNAMODB_TTL_DEFAULT_VALUE

from tests.sample_data.sample_data import ORGANIZATION_ID

DATA_PREFIX = 'datum'
CURRENT_TIME = '2023-12-31T23:59:59+00:00'


@pytest.fixture(scope='function')
def add_items(mock_table, count=20):
    for x in range(0, count):
        mock_table.put_item(Item={
            'PK': f'organization#{ORGANIZATION_ID}',
            'SK': f'{DATA_PREFIX}#{x}'})


@pytest.fixture(scope='function')
def context(mocker, add_items):
    import czc.dynamodb as handler
    with Context(mocker, handler, [], extra={'handler': handler}) as context:
        yield context


@pytest.mark.unit
@pytest.mark.parametrize('parts, composed_key', [
    (['one'], 'one'),
    (['one', 'two'], 'one#two'),
    (['one', '2', 'three'], 'one#2#three')
])
def test_compose_and_decompose_key(context, parts, composed_key):
    assert context.handler.compose_key(*parts) == composed_key
    assert context.handler.decompose_key(composed_key) == parts


@pytest.mark.unit
@pytest.mark.parametrize('parts, composed_key, composite_char', [
    (['one'], 'one', '-'),
    (['one', 'two'], 'one-two', '-'),
    (['one', '2', 'three'], 'one/2/three', '/')
])
def test_compose_and_decompose_key_with_composite_char(context, parts, composed_key, composite_char):
    assert context.handler.compose_key(*parts, composite_char=composite_char) == composed_key
    assert context.handler.decompose_key(composed_key, composite_char=composite_char) == parts


@pytest.mark.unit
@pytest.mark.parametrize('prefix, key, expected_key', [
    ('prefix', 'key', 'prefix#key'),
    ('prefix', 'prefix-key', 'prefix#prefix-key'),
    ('prefix', 'prefix#key', 'prefix#key'),
])
def test_prefix(context, prefix, key, expected_key):
    assert context.handler.prefix(prefix, key) == expected_key


@pytest.mark.unit
@pytest.mark.parametrize('prefix, key, expected_key, composite_char', [
    ('prefix', 'key', 'prefix-key', '-'),
    ('prefix', 'prefix-key', 'prefix/prefix-key', '/'),
    ('prefix', 'prefix#key', 'prefix*prefix#key', '*'),
    ('prefix', 'prefix*key', 'prefix*key', '*'),
])
def test_prefix_with_composite_char(context, prefix, key, expected_key, composite_char):
    assert context.handler.prefix(prefix, key, composite_char=composite_char) == expected_key


@pytest.mark.unit
@pytest.mark.parametrize('prefix, key, expected_key', [
    ('prefix', 'key', 'key'),
    ('prefix', 'prefix#key', 'key'),
    ('prefix', 'prefix-key', 'prefix-key'),
    ('prefix', 'prefix#prefix-key', 'prefix-key'),
])
def test_unfix(context, prefix, key, expected_key):
    assert context.handler.unfix(prefix, key) == expected_key


@pytest.mark.unit
@pytest.mark.parametrize('prefix, key, expected_key, composite_char', [
    ('prefix', 'key', 'key', '-'),
    ('prefix', 'prefix#key', 'prefix#key', '-'),
    ('prefix', 'prefix-key', 'key', '-'),
    ('prefix', 'prefix-prefix#key', 'prefix#key', '-'),
])
def test_unfix_with_composite_char(context, prefix, key, expected_key, composite_char):
    assert context.handler.unfix(prefix, key, composite_char=composite_char) == expected_key


@pytest.mark.unit
@pytest.mark.parametrize('scan_kwargs', [
    pytest.param({'FilterExpression': Attr('SK').begins_with(DATA_PREFIX),
                  'ProjectionExpression': 'SK,cz_organization_id'}, id='project_expression'),
    pytest.param({}, id='empty_kwargs'),
])
def test_scan(context, mock_table, add_items, scan_kwargs):
    items = context.handler.scan(mock_table, scan_kwargs)
    assert len(items) > 0


@pytest.mark.unit
@pytest.mark.parametrize('query_kwargs', [
    pytest.param({
        'KeyConditionExpression': (Key('PK').eq(
            f'organization#{ORGANIZATION_ID}') & Key('SK').begins_with(DATA_PREFIX)),
        'ScanIndexForward': False}, id='get')
])
def test_query(context, mock_table, add_items, query_kwargs):
    items = context.handler.query(mock_table, query_kwargs)
    assert len(items) > 0


def _set_ttl_info(environ, has_test_key, ttl_override_key, ttl_override_value):
    if has_test_key:
        environ['cloudzero_test_key'] = 'test-key'
    elif 'cloudzero_test_key' in environ:
        del environ['cloudzero_test_key']
    if ttl_override_key is not None:
        environ['DYNAMODB_TTL_OVERRIDE'] = ttl_override_key
    elif 'DYNAMODB_TTL_OVERRIDE' in environ:
        del environ['DYNAMODB_TTL_OVERRIDE']
    if ttl_override_value is not None:
        environ['DYNAMODB_TTL_OVERRIDE_VALUE'] = ttl_override_value
    elif 'DYNAMODB_TTL_OVERRIDE_VALUE' in environ:
        del environ['DYNAMODB_TTL_OVERRIDE_VALUE']


def _get_expected_expiry_time(ttl_value):
    return int((datetime.fromisoformat(CURRENT_TIME) + timedelta(seconds=ttl_value)).timestamp())


@pytest.mark.unit
@pytest.mark.parametrize('item', [
    ({'x': 1, 'y': 'foo', 'z': 'bar'}),
    ({'some': 'data', 'goes': 'here'}),
])
@pytest.mark.parametrize('has_test_key,ttl_override_key,ttl_override_value,expected_ttl_info', [
    (False, 'xyz', None, {}),
    (True, None, None, {'ttl': _get_expected_expiry_time(DYNAMODB_TTL_DEFAULT_VALUE)}),
    (True, 'xyz', None, {'xyz': _get_expected_expiry_time(DYNAMODB_TTL_DEFAULT_VALUE)}),
    (True, None, '300', {'ttl': _get_expected_expiry_time(300)}),
    (True, 'xyz', '500', {'xyz': _get_expected_expiry_time(500)}),
])
@pytest.mark.parametrize('has_condition', [True, False])
@pytest.mark.parametrize('has_return_values', [True, False])
@freeze_time(CURRENT_TIME)
def test_put_item(context, environ, item, has_test_key, ttl_override_key, ttl_override_value, expected_ttl_info,
                  has_condition, has_return_values):
    _set_ttl_info(environ, has_test_key, ttl_override_key, ttl_override_value)

    table = Mock()
    current_item = {'Item': 'data'}
    table.put_item.return_value = current_item

    condition_args = dict(ConditionExpression=Mock()) if has_condition else {}
    return_values_args = dict(ReturnValues=Mock()) if has_return_values else {}
    result = context.handler.put_item(table, Item=item, **condition_args, **return_values_args)
    expected_item = {**item, **expected_ttl_info}
    table.put_item.assert_called_once_with(Item=expected_item, **condition_args, **return_values_args)
    assert result == current_item


@pytest.mark.unit
def test_atomic_update(context, environ, mock_table):
    key = {
        'PK': f'organization#{ORGANIZATION_ID}',
        'SK': f'{DATA_PREFIX}#update_test'
    }

    mock_table.put_item(Item={
        **key,
        'field': int(5)
    })
    failures = 2

    @context.handler.atomic()
    def update():
        item = mock_table.get_item(Key=key)['Item']
        nonlocal failures

        if failures > 0:
            mock_table.put_item(Item=update_in(item, ['field'], lambda x: int(x + 1)))
            failures -= 1

        return context.handler.update_item_from_dict(mock_table, key, {'field': int(item['field'] + 1)}, ConditionExpression=Attr('field').eq(item['field']))

    update()
    new_item = mock_table.get_item(Key=key)['Item']

    assert int(new_item['field']) == 8


class TestStrEnum(str, Enum):
    value1 = 'str_value1'
    value2 = 'str_value2'


class TestIntEnum(int, Enum):
    value1 = 1
    value2 = 2


@pytest.mark.unit
@pytest.mark.parametrize('use_enum_values', [True, False])
@pytest.mark.parametrize('enum_to_test', [*TestStrEnum, *TestIntEnum])
def test_update_item_from_dict_handles_enums_correctly(context, environ, mock_table, use_enum_values, enum_to_test):
    expected_result = enum_to_test.value if use_enum_values else str(enum_to_test)

    key = {
        'PK': f'organization#{ORGANIZATION_ID}',
        'SK': f'{DATA_PREFIX}#update_test'
    }

    response = context.handler.update_item_from_dict(mock_table, key, {'enum_field': enum_to_test}, use_enum_values=use_enum_values)
    assert response['enum_field'] == expected_result

    new_item = mock_table.get_item(Key=key)['Item']

    assert new_item['enum_field'] == expected_result


@pytest.mark.unit
@pytest.mark.parametrize('has_test_key,ttl_override_key,ttl_override_value,expected_ttl_info', [
    (False, 'xyz', None, {}),
    (True, None, None, {'ttl': _get_expected_expiry_time(DYNAMODB_TTL_DEFAULT_VALUE)}),
    (True, 'xyz', None, {'xyz': _get_expected_expiry_time(DYNAMODB_TTL_DEFAULT_VALUE)}),
    (True, None, '300', {'ttl': _get_expected_expiry_time(300)}),
    (True, 'xyz', '500', {'xyz': _get_expected_expiry_time(500)}),
])
@freeze_time(CURRENT_TIME)
def test_update_item_from_dict_test_key(context, environ, mock_table, has_test_key, ttl_override_key, ttl_override_value, expected_ttl_info):
    _set_ttl_info(environ, has_test_key, ttl_override_key, ttl_override_value)

    key = {
        'PK': f'organization#{ORGANIZATION_ID}',
        'SK': f'{DATA_PREFIX}#update_test'
    }
    item = {'x': 1, 'y': 'foo', 'z': 'bar'}

    context.handler.update_item_from_dict(mock_table, key=key, dictionary=item)
    new_item = mock_table.get_item(Key=key)['Item']
    if has_test_key:
        ttl_key = ttl_override_key or 'ttl'
        assert new_item[ttl_key] == Decimal(expected_ttl_info.get(ttl_key, DYNAMODB_TTL_DEFAULT_VALUE))
    else:
        assert ttl_override_key not in new_item
