# 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
from typing import NamedTuple, List
from unittest.mock import Mock, call
from uuid import uuid4

import pytest
from czt.unittest import Context, MockEx, MATCH

from czc import feature_data
from czc.feature_data import FeatureDataConnection
from czc.type_util import deserializable


@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, ['get_object', 'list_objects', 'list_all_objects',
                                        'set_object', 'update_object_atomic', 'delete_object'], extra=extra) as context:
        yield context


@deserializable()
class DataObject(NamedTuple):
    val: int
    name: str
    data: List[str]


TEST_DATA = [
    (dict(val=1, name='foo', data=['one', 'two']), DataObject(val=1, name='foo', data=['one', 'two'])),
    (dict(val=2, name='bar', data=['three', 'four']), DataObject(val=2, name='bar', data=['three', 'four']))
]

TEST_DATA_LISTS = [
    ([x[0] for x in TEST_DATA[:n]], [x[1] for x in TEST_DATA[:n]]) for n in range(3)
]


@pytest.mark.unit
@pytest.mark.parametrize('data,expected_result', [
    ({}, None),
    *TEST_DATA
])
def test_get_should_coerce_result(context, data, expected_result):
    args = (context.conn, context.org_id, context.keyspace, 'k1', 'k2')
    MockEx(context.mock_get_object).if_called_with(*args).return_value(data)

    result = feature_data.for_type(DataObject).get(*args)

    assert result == expected_result


@pytest.mark.unit
@pytest.mark.parametrize('data,expected_result', TEST_DATA_LISTS)
def test_list_should_coerce_all_results(context, data, expected_result):
    args = (context.conn, context.org_id, context.keyspace, 'k1', 'k2')
    MockEx(context.mock_list_objects).if_called_with(*args).return_value(data)

    result = list(feature_data.for_type(DataObject).list(*args))

    assert result == expected_result


@pytest.mark.unit
@pytest.mark.parametrize('data,expected_result', TEST_DATA_LISTS)
def test_list_all_should_coerce_all_results(context, data, expected_result):
    args = (context.conn, context.keyspace, 'k1', 'k2')
    MockEx(context.mock_list_all_objects).if_called_with(*args).return_value(data)

    result = list(feature_data.for_type(DataObject).list_all(*args))

    assert result == expected_result


@pytest.mark.unit
@pytest.mark.parametrize('serializable_data,object_data', TEST_DATA)
def test_set_should_call_set_object_with_serializable_data(context, serializable_data, object_data):
    args = (context.conn, context.org_id, context.keyspace, 'k1', 'k2')

    feature_data.for_type(DataObject).set(*args, data=object_data)

    context.mock_set_object.assert_called_once_with(*args, data=serializable_data)


@pytest.mark.unit
@pytest.mark.parametrize('serializable_data,expected_object_data', TEST_DATA)
def test_update_should_use_typed_handler(context, serializable_data, expected_object_data):
    args = (context.conn, context.org_id, context.keyspace, 'k1', 'k2')

    context.mock_update_object_atomic.return_value = serializable_data

    new_val = 99
    new_object_data = expected_object_data._replace(val=new_val)
    update_handler = Mock()
    update_handler.return_value = new_object_data
    result = feature_data.for_type(DataObject).update(*args, handler=update_handler)

    def validate_handler(h):
        result = h(serializable_data)
        return (update_handler.mock_calls == [call(expected_object_data)] and
                result == {**serializable_data, 'val': new_val})

    context.mock_update_object_atomic.assert_called_once_with(*args, update_handler=MATCH(validate_handler))
    assert result == expected_object_data


@pytest.mark.unit
@pytest.mark.parametrize('serializable_data,expected_handler_param', [
    ({}, None)
])
def test_update_should_handle_empty_object(context, serializable_data, expected_handler_param):
    args = (context.conn, context.org_id, context.keyspace, 'k1', 'k2')

    context.mock_update_object_atomic.return_value = None

    update_handler = Mock()
    update_handler.side_effect = lambda x: x
    result = feature_data.for_type(DataObject).update(*args, handler=update_handler)

    def validate_handler(h):
        h(serializable_data)
        return update_handler.mock_calls == [call(expected_handler_param)]

    context.mock_update_object_atomic.assert_called_once_with(*args, update_handler=MATCH(validate_handler))
    assert result is None


@pytest.mark.unit
def test_delete_should_passthrough_args(context):
    args = (context.conn, context.keyspace, 'k1', 'k2')

    feature_data.for_type(DataObject).delete(*args)

    context.mock_delete_object.assert_called_once_with(*args)
