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

import pytest
from pydantic import BaseModel

from czc import feature_data
from czc.feature_data import FeatureDataConnection
from czc.type_util import deserializable
from czt.unittest import MATCH, Context, MockEx

UTC = timezone.utc


@pytest.fixture(scope='function')
def context(mocker):
    extra = dict(
        conn=FeatureDataConnection(_table=Mock(), _s3_bucket=str(uuid4()), _ttl=None, _s3_temp_prefix='tmp'),
        org_id=str(uuid4()),
        keyspace=str(uuid4()),
    )
    with Context(
        mocker,
        feature_data,
        [
            'get_object',
            'list_objects',
            'list_objects_with_keys',
            'list_all_objects',
            'list_all_objects_with_organization',
            'set_object',
            'update_object_atomic',
            'delete_object',
        ],
        extra=extra,
    ) as context:
        yield context


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


class Enum1(Enum):
    A = 1
    B = 2
    C = 3


@deserializable()
class Model1(BaseModel):
    a: int
    b: float
    c: str
    d: dict[str, str]
    e: list[str]
    f: datetime
    g: Enum1
    h: DataObject
    i: tuple[str, str]


@deserializable()
class Model2(BaseModel):
    a: str
    b: 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)]

TEST_DATA_LISTS_WITH_ORG = [
    ([(org_id, x[0]) for x in TEST_DATA[:n]], [(org_id, x[1]) for x in TEST_DATA[:n]])
    for n in range(3)
    for org_id in (str(uuid4()),)
]


@pytest.mark.unit
@pytest.mark.parametrize(
    'record',
    [
        Model1(
            a=1,
            b=1.2,
            c='abc',
            d={'x': 'X', 'y': 'Y'},
            e=['a', 'b'],
            f=datetime(2020, 3, 14, 17, 54, 57, 415748, UTC),
            g=Enum1.B,
            h=DataObject(val=1, name='foo', data=['one', 'two']),
            i=('three', 'four'),
        ),
    ],
)
def test_roundtrip(mock_table, record):
    conn = FeatureDataConnection(_table=mock_table, _s3_bucket=str(uuid4()), _ttl=None, _s3_temp_prefix='tmp')
    org_id = str(uuid4())
    keyspace = str(uuid4())
    key = ('k1', 'k2')
    f = feature_data.for_type(Model1)
    f.set(conn, org_id, keyspace, *key, data=record)
    assert f.get(conn, org_id, keyspace, *key) == record


@pytest.mark.unit
@pytest.mark.parametrize('record', [Model2(a='x', b='')])
def test_roundtrip_strings(mock_table, record):
    # emptry strings were previously replaced with None, causing an error upon retrieval
    # when using pydantic models
    conn = FeatureDataConnection(_table=mock_table, _s3_bucket=str(uuid4()), _ttl=None, _s3_temp_prefix='tmp')
    org_id = str(uuid4())
    keyspace = str(uuid4())
    key = ('k1', 'k2')
    f = feature_data.for_type(Model2)
    f.set(conn, org_id, keyspace, *key, data=record)
    assert f.get(conn, org_id, keyspace, *key) == record


@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('data,expected_result', TEST_DATA_LISTS_WITH_ORG)
def test_list_all_with_organization_should_coerce_all_results(context, data, expected_result):
    args = (context.conn, context.keyspace, 'k1', 'k2')
    MockEx(context.mock_list_all_objects_with_organization).if_called_with(*args).return_value(data)

    result = list(feature_data.for_type(DataObject).list_all_with_organization(*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)


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


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

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

    assert result == expected_result
