# 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 os
from decimal import Decimal
from unittest.mock import ANY, Mock, call
from uuid import uuid4

import pytest
from boto3.dynamodb.conditions import Key, Attr
from botocore.exceptions import ClientError
from czt.unittest import Context, MockEx

from czc import dynamodb as dyn
from czc import feature_data
from czc.feature_data import FeatureDataConnection, S3_OBJECT_STORE_ROOT

EXCEED_MAX_SIZE_ERROR = ClientError(
    error_response=dict(
        Error=dict(
            Code='ValidationException',
            Message='Item size has exceeded the maximum allowed size')
    ), operation_name='PutItem')


@pytest.fixture(scope='function')
def context(mocker):
    extra = dict(conn=FeatureDataConnection(_table=Mock(), _s3_bucket=str(uuid4())),
                 org_id=str(uuid4()),
                 keyspace=str(uuid4()))
    with Context(mocker, feature_data, ['dyn', 's3', 'uuid', 's3_client'], extra=extra) as context:
        os.environ['NAMESPACE'] = 'ns'
        context.mock_dyn.prefix.side_effect = lambda p, x: f'{p}#{x}'
        context.mock_dyn.compose_key.side_effect = lambda *x: '#'.join(x)
        context.mock_dyn.type_for_dynamo.side_effect = dyn.type_for_dynamo
        context.mock_uuid.uuid4.return_value = 'UUID'
        mocker.patch('retrying.time.sleep')

        yield context


@pytest.mark.unit
@pytest.mark.parametrize('table_name,bucket_name', [
    ('standard-table', 'standard-bucket')
])
def test_connect_should_use_standard_table_and_bucket_names(context, table_name, bucket_name):
    expected_table = Mock()
    MockEx(context.mock_s3.get_bucket_name).if_called_with().return_value(bucket_name)
    MockEx(context.mock_dyn.get_table_name).if_called_with().return_value(table_name)
    MockEx(context.mock_dyn.table).if_called_with(table_name).return_value(expected_table)

    conn = feature_data.connect()
    assert conn._table == expected_table
    assert conn._s3_bucket == bucket_name


@pytest.mark.unit
@pytest.mark.parametrize('org_id,keyspace,keys,expected_dynamo_key', [
    ('org1', 'ks', [], dict(PK='organization#org1', SK='ks')),
    ('org2', 'obj', ['k1', 'k2'], dict(PK='organization#org2', SK='obj#k1#k2'))
])
def test_get_object_should_use_correct_dynamo_key(context, org_id, keyspace, keys, expected_dynamo_key):
    feature_data.get_object(context.conn, org_id, keyspace, *keys)

    context.conn._table.get_item.assert_called_once_with(Key=expected_dynamo_key, ConsistentRead=True)


@pytest.mark.unit
@pytest.mark.parametrize('dynamo_data,expected_result', [
    (None, {}),
    ({}, {}),
    ({'val': 1}, {'val': 1})
])
def test_get_object_should_return_dynamo_data(context, dynamo_data, expected_result):
    MockEx(context.conn._table.get_item).if_called_with_anything().return_value(
        dict(Item=dict(data=dynamo_data)) if dynamo_data is not None else dict()
    )
    result = feature_data.get_object(context.conn, context.org_id, context.keyspace)

    assert result == expected_result


@pytest.mark.unit
@pytest.mark.parametrize('dynamo_data,expected_path,expected_result', [
    ('REF:path/to/data', 'path/to/data', {'data': 1}),
    ('Invalid string', '', ValueError)
])
def test_get_object_should_return_s3_data_for_ref(context, dynamo_data, expected_path, expected_result):
    MockEx(context.conn._table.get_item).if_called_with_anything().return_value(
        dict(Item=dict(data=dynamo_data)) if dynamo_data is not None else dict()
    )
    MockEx(context.mock_s3.get_json_object).if_called_with(
        context.conn._s3_bucket, f'{S3_OBJECT_STORE_ROOT}/{expected_path}').return_value(expected_result)

    def get_object():
        return feature_data.get_object(context.conn, context.org_id, context.keyspace)

    if expected_result == ValueError:
        pytest.raises(ValueError, get_object)
    else:
        assert get_object() == expected_result


@pytest.mark.unit
@pytest.mark.parametrize('org_id,keyspace,keys,expected_dynamo_condition', [
    ('org1', 'ks', [], Key('PK').eq('organization#org1') & Key('SK').begins_with('ks#')),
    ('org2', 'obj', ['k1', 'k2'], Key('PK').eq('organization#org2') & Key('SK').begins_with('obj#k1#k2#'))
])
def test_list_objects_should_use_correct_dynamo_condition(context, org_id, keyspace, keys, expected_dynamo_condition):
    list(feature_data.list_objects(context.conn, org_id, keyspace, *keys))

    context.mock_dyn.query.assert_called_once_with(
        context.conn._table, dict(KeyConditionExpression=expected_dynamo_condition, ConsistentRead=True))


@pytest.mark.unit
@pytest.mark.parametrize('dynamo_data,expected_result', [
    ([], []),
    ([{'val': 1}, {'val': 2}], [{'val': 1}, {'val': 2}])
])
def test_list_objects_should_return_dynamo_data(context, dynamo_data, expected_result):
    MockEx(context.mock_dyn.query).if_called_with_anything().return_value(
        [dict(data=x) for x in dynamo_data]
    )
    result = list(feature_data.list_objects(context.conn, context.org_id, context.keyspace))

    assert result == expected_result


@pytest.mark.unit
@pytest.mark.parametrize('dynamo_data,expected_s3_data,expected_result', [
    (['REF:path/to/data'], {'path/to/data': {'data': 1}}, [{'data': 1}]),
    (['REF:path/to/data', 'REF:path2'],
     {'path/to/data': {'data': 1}, 'path2': {'data': 2}}, [{'data': 1}, {'data': 2}]),
    (['REF:path/to/data', {'data': 2}], {'path/to/data': {'data': 1}}, [{'data': 1}, {'data': 2}]),
    (['Invalid string'], {}, ValueError)
])
def test_list_objects_should_return_s3_data_for_ref(context, dynamo_data, expected_s3_data, expected_result):
    MockEx(context.mock_dyn.query).if_called_with_anything().return_value(
        [dict(data=x) for x in dynamo_data]
    )
    for path, data in expected_s3_data.items():
        MockEx(context.mock_s3.get_json_object).if_called_with(
            context.conn._s3_bucket, f'{S3_OBJECT_STORE_ROOT}/{path}').return_value(data)

    def list_objects():
        return list(feature_data.list_objects(context.conn, context.org_id, context.keyspace))

    if expected_result == ValueError:
        pytest.raises(ValueError, list_objects)
    else:
        assert list_objects() == expected_result


@pytest.mark.unit
@pytest.mark.parametrize('keyspace,keys,expected_dynamo_condition', [
    ('ks', [], Attr('PK').begins_with('organization#') & Attr('SK').begins_with('ks#')),
    ('obj', ['k1', 'k2'], Attr('PK').begins_with('organization#') & Attr('SK').begins_with('obj#k1#k2#'))
])
def test_list_all_objects_should_use_correct_dynamo_condition(context, keyspace, keys, expected_dynamo_condition):
    list(feature_data.list_all_objects(context.conn, keyspace, *keys))

    context.mock_dyn.scan.assert_called_once_with(context.conn._table,
                                                  dict(FilterExpression=expected_dynamo_condition))


@pytest.mark.unit
@pytest.mark.parametrize('dynamo_data,expected_result', [
    ([], []),
    ([{'val': 1}, {'val': 2}], [{'val': 1}, {'val': 2}])
])
def test_list_all_objects_should_return_dynamo_data(context, dynamo_data, expected_result):
    MockEx(context.mock_dyn.scan).if_called_with_anything().return_value(
        [dict(data=x) for x in dynamo_data]
    )
    result = list(feature_data.list_all_objects(context.conn, context.keyspace))

    assert result == expected_result


@pytest.mark.unit
@pytest.mark.parametrize('dynamo_data,expected_s3_data,expected_result', [
    (['REF:path/to/data'], {'path/to/data': {'data': 1}}, [{'data': 1}]),
    (['REF:path/to/data', 'REF:path2'],
     {'path/to/data': {'data': 1}, 'path2': {'data': 2}}, [{'data': 1}, {'data': 2}]),
    (['REF:path/to/data', {'data': 2}], {'path/to/data': {'data': 1}}, [{'data': 1}, {'data': 2}]),
    (['Invalid string'], {}, ValueError)
])
def test_list_all_objects_should_return_s3_data_for_ref(context, dynamo_data, expected_s3_data, expected_result):
    MockEx(context.mock_dyn.scan).if_called_with_anything().return_value(
        [dict(data=x) for x in dynamo_data]
    )
    for path, data in expected_s3_data.items():
        MockEx(context.mock_s3.get_json_object).if_called_with(
            context.conn._s3_bucket, f'{S3_OBJECT_STORE_ROOT}/{path}').return_value(data)

    def list_objects():
        return list(feature_data.list_all_objects(context.conn, context.keyspace))

    if expected_result == ValueError:
        pytest.raises(ValueError, list_objects)
    else:
        assert list_objects() == expected_result


@pytest.mark.unit
@pytest.mark.parametrize('org_id,keyspace,keys,expected_dynamo_key', [
    ('org1', 'ks', [], dict(PK='organization#org1', SK='ks')),
    ('org2', 'obj', ['k1', 'k2'], dict(PK='organization#org2', SK='obj#k1#k2'))
])
def test_set_object_should_use_correct_dynamo_key(context, org_id, keyspace, keys, expected_dynamo_key):
    feature_data.set_object(context.conn, org_id, keyspace, *keys, data={})

    context.mock_dyn.put_item.assert_called_once_with(context.conn._table,
                                                      ReturnValues='ALL_OLD',
                                                      Item=dict(**expected_dynamo_key, data=ANY))


@pytest.mark.unit
@pytest.mark.parametrize('data,expected_put_data', [
    ({}, {}),
    ({'val': 1}, {'val': 1}),
    ({'val': 1.23}, {'val': Decimal('1.23')})
])
def test_set_object_should_put_dynamo_data(context, data, expected_put_data):
    feature_data.set_object(context.conn, context.org_id, context.keyspace, data=data)

    context.mock_dyn.put_item.assert_called_once_with(context.conn._table,
                                                      ReturnValues='ALL_OLD',
                                                      Item=dict(PK=ANY, SK=ANY, data=expected_put_data))


@pytest.mark.unit
@pytest.mark.parametrize('org_id,keyspace,keys,expected_ref', [
    ('org1', 'ks', [], 'organization_org1/ks_UUID.json'),
    ('aaa-bbb', 'obj', ['key@1', 'key@2'], 'organization_aaa_bbb/obj_key_1_key_2_UUID.json')
])
def test_set_object_should_put_large_data_in_s3_with_ref(context, org_id, keyspace, keys, expected_ref):
    context.mock_dyn.put_item.side_effect = [EXCEED_MAX_SIZE_ERROR, None]
    feature_data.set_object(context.conn, org_id, keyspace, *keys, data={})

    context.mock_dyn.put_item.assert_called_with(context.conn._table,
                                                 ReturnValues='ALL_OLD',
                                                 Item=dict(PK=ANY, SK=ANY, data=f'REF:{expected_ref}'))
    context.mock_s3.put_json_object.assert_called_once_with(context.conn._s3_bucket,
                                                            f'{S3_OBJECT_STORE_ROOT}/{expected_ref}',
                                                            ANY)


@pytest.mark.unit
@pytest.mark.parametrize('data,expected_s3_data', [
    ({}, {}),
    ({'val': 1}, {'val': 1})
])
def test_set_object_should_put_large_data_in_s3(context, data, expected_s3_data):
    context.mock_dyn.put_item.side_effect = [EXCEED_MAX_SIZE_ERROR, None]
    feature_data.set_object(context.conn, context.org_id, context.keyspace, data=data)

    context.mock_s3.put_json_object.assert_called_once_with(context.conn._s3_bucket,
                                                            ANY,
                                                            expected_s3_data)


@pytest.mark.parametrize('prior_data,expected_delete_key', [
    ({'val': 1}, None),
    ('REF:key.json', 'key.json')
])
def test_set_object_should_delete_prior_data_from_s3(context, prior_data, expected_delete_key):
    context.mock_dyn.put_item.return_value = dict(
        Attributes=dict(data=prior_data)
    )
    feature_data.set_object(context.conn, context.org_id, context.keyspace, data={})

    if expected_delete_key is not None:
        context.mock_s3_client.delete_object.assert_called_once_with(
            Bucket=context.conn._s3_bucket,
            Key=f'{S3_OBJECT_STORE_ROOT}/{expected_delete_key}')
    else:
        assert not context.mock_s3_client.delete_object.called


@pytest.mark.unit
@pytest.mark.parametrize('current_data,expected_object', [
    ({}, {}),
    ({'data': None}, {}),
    ({'data': {}}, {}),
    ({'data': {'key': 'value'}}, {'key': 'value'}),
    ({'data': 'BAD DATA'}, ValueError),
])
def test_update_object_data_atomic_should_call_handler_with_current_data(context, current_data, expected_object):
    MockEx(context.conn._table.get_item).if_called_with_anything().return_value({'Item': current_data})

    mock_update_function = Mock()

    def update():
        feature_data.update_object_atomic(context.conn, context.org_id, context.keyspace,
                                          update_handler=mock_update_function)

    if expected_object is ValueError:
        pytest.raises(expected_object, update)
    else:
        update()
        mock_update_function.assert_called_once_with(expected_object)


@pytest.mark.unit
@pytest.mark.parametrize('new_data,should_put_item', [
    (None, False),
    ({'key': 'value'}, True),
    ({'key': 'value', 'k2': 123}, True)
])
def test_update_object_data_atomic_should_put_new_data_with_handler_result(context, new_data, should_put_item):
    mock_update_function = Mock()
    mock_update_function.return_value = new_data

    feature_data.update_object_atomic(context.conn, context.org_id, context.keyspace,
                                      update_handler=mock_update_function)

    assert context.mock_dyn.put_item.called == should_put_item
    if should_put_item:
        context.mock_dyn.put_item.assert_called_with(context.conn._table,
                                                     ConditionExpression=ANY,
                                                     ReturnValues='ALL_OLD',
                                                     Item=dict(PK=ANY, SK=ANY, data=new_data))


@pytest.mark.unit
@pytest.mark.parametrize('exception,should_retry', [
    (ClientError(error_response=dict(Error=dict(Code='ConditionalCheckFailedException')),
                 operation_name='PutItem'), True),
    (ValueError('other error'), False),
    (ClientError(error_response=dict(Error=dict(Code='ValidationException')),
                 operation_name='PutItem'), False),
])
def test_update_object_data_atomic_should_retry_if_condition_check_fails(context, exception, should_retry):
    new_data = {'test': 'data'}
    mock_update_function = Mock()
    mock_update_function.return_value = new_data
    context.mock_dyn.put_item.side_effect = [exception, None]

    def update():
        feature_data.update_object_atomic(context.conn, context.org_id, context.keyspace,
                                          update_handler=mock_update_function)

    expected_call = call(context.conn._table,
                         ConditionExpression=ANY,
                         ReturnValues='ALL_OLD',
                         Item=dict(PK=ANY, SK=ANY, data=new_data))
    if should_retry:
        update()
        assert context.mock_dyn.put_item.mock_calls == [expected_call, expected_call]
    else:
        pytest.raises(type(exception), update)


@pytest.mark.unit
@pytest.mark.parametrize('org_id,keyspace,keys,expected_dynamo_key', [
    ('org1', 'ks', [], dict(PK='organization#org1', SK='ks')),
    ('org2', 'obj', ['k1', 'k2'], dict(PK='organization#org2', SK='obj#k1#k2'))
])
def test_delete_object_should_delete_dynamo_data(context, org_id, keyspace, keys, expected_dynamo_key):
    feature_data.delete_object(context.conn, org_id, keyspace, *keys)

    context.conn._table.delete_item.assert_called_once_with(
        ReturnValues='ALL_OLD',
        Key=expected_dynamo_key)


@pytest.mark.unit
@pytest.mark.parametrize('prior_data,expected_delete_key', [
    ({'val': 1}, None),
    ('REF:key.json', 'key.json')
])
def test_delete_object_should_delete_prior_data_from_s3(context, prior_data, expected_delete_key):
    context.conn._table.delete_item.return_value = dict(
        Attributes=dict(data=prior_data)
    )
    feature_data.delete_object(context.conn, context.org_id, context.keyspace)

    if expected_delete_key is not None:
        context.mock_s3_client.delete_object.assert_called_once_with(
            Bucket=context.conn._s3_bucket,
            Key=f'{S3_OBJECT_STORE_ROOT}/{expected_delete_key}')
    else:
        assert not context.mock_s3_client.delete_object.called
