# Copyright (c) 2024 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 enum import Enum
from typing import Any, Dict, List, NamedTuple, Set
from unittest.mock import ANY, Mock, call
from uuid import uuid4

import cachetools
import pytest
from boto3.dynamodb.conditions import ConditionExpressionBuilder, Key

from czc import dynamodb as dyn
from czc import fast_distributed_cache
from czc.fast_distributed_cache import (
    KEYSPACE,
    ORGANIZATION_KEY,
    ORGANIZATION_NONE,
    PERSISTENCE_TTL,
    RESULT_KEY,
    TOKEN_KEY,
    _CachedResult,
    cached,
    invalidate,
)
from czc.type_util import deserializable
from czt.unittest import MATCH, Context


@deserializable()
class TestResult(NamedTuple):
    name: str
    value: int


@pytest.fixture(scope='function')
def lrucache_spy(mocker):
    spy = mocker.spy(cachetools, 'LRUCache')
    yield spy


@pytest.fixture(scope='function')
def context(mocker):
    with Context(
        mocker,
        fast_distributed_cache,
        [
            'dyn',
            'feature_data',
            'uuid4',
            'result_store',
        ],
        extra={'mock_table': Mock(), 'mock_conn': Mock()},
    ) as context:
        context.mock_dyn.table.return_value = context.mock_table
        fast_distributed_cache._dynamo_table.cache_clear()

        context.mock_feature_data.connect.return_value = context.mock_conn
        context.mock_uuid4.return_value = 'test-uuid'

        context.mock_dyn.COMPOSITE_CHAR = dyn.COMPOSITE_CHAR
        context.mock_dyn.prefix = dyn.prefix
        context.mock_dyn.compose_key = dyn.compose_key
        context.mock_dyn.decompose_key = dyn.decompose_key

        context.mock_result_store.get.return_value = None
        context.mock_table.get_item.return_value = {'Item': {'data': 'token-123'}}

        yield context


@pytest.mark.unit
@pytest.mark.parametrize(
    'current_token,cached_data,expected_result,expected_set_called',
    [
        pytest.param('token-123', {'key': 'value'}, {'key': 'value'}, False, id='token_match_returns_cached_result'),
        pytest.param('different-token', {'old': 'data'}, {'computed': 'value'}, True, id='token_mismatch_calls_original'),
        pytest.param(None, None, {'computed': 'value'}, True, id='no_cached_result_calls_original'),
    ],
)
def test_cached_function_return_value_based_on_token_and_cache(
    context, current_token: str, cached_data, expected_result, expected_set_called
):
    cached_token = 'token-123'
    cached_result = _CachedResult(token=cached_token, parameters=None, data=cached_data) if cached_data is not None else None
    context.mock_result_store.get.return_value = cached_result

    token_data = {'Item': {'data': current_token}} if current_token else {}
    context.mock_table.get_item.return_value = token_data

    @cached('test-cache')
    def test_function(param1: str, param2: int) -> Dict[str, Any]:
        return {'computed': 'value'}

    result = test_function('test', 123)

    assert result == expected_result
    context.mock_feature_data.connect.assert_called_once_with(ttl=PERSISTENCE_TTL)

    expected_set_token = current_token or 'test-uuid'
    expected_set_calls = (
        [
            call(
                context.mock_conn,
                ORGANIZATION_NONE,
                KEYSPACE,
                'test-cache',
                ANY,
                RESULT_KEY,
                data=_CachedResult(token=expected_set_token, parameters=None, data={'computed': 'value'}),
            )
        ]
        if expected_set_called
        else []
    )
    assert context.mock_result_store.set.call_args_list == expected_set_calls


@pytest.mark.unit
@pytest.mark.parametrize(
    'save_params,all_params,expected_stored_params',
    [
        pytest.param(
            {'param1', 'param2'},
            {'param1': 'test', 'param2': 123, 'param3': True},
            {'param1': 'test', 'param2': 123},
            id='save_specific_params',
        ),
        pytest.param(
            '*',
            {'param1': 'test', 'param2': 123, 'param3': True},
            {'param1': 'test', 'param2': 123, 'param3': True},
            id='save_all_params',
        ),
        pytest.param(None, {'param1': 'test', 'param2': 123, 'param3': True}, None, id='save_no_params'),
    ],
)
def test_cached_function_saves_parameters_correctly(context, save_params: Set[str], all_params: dict, expected_stored_params):
    @cached('test-cache', save_params=save_params)
    def test_function(param1: str, param2: int, param3: bool = False) -> Dict[str, Any]:
        return {'computed': 'value'}

    test_function(**all_params)

    context.mock_result_store.set.assert_called_once_with(
        context.mock_conn,
        ORGANIZATION_NONE,
        KEYSPACE,
        'test-cache',
        ANY,
        RESULT_KEY,
        data=MATCH(lambda x: x.parameters == expected_stored_params),
    )


@pytest.mark.unit
@pytest.mark.parametrize(
    'param_name,org_id,expected_sk_org_id',
    [
        pytest.param('organization_id', 'org-1', 'org-1', id='organization_id_param'),
        pytest.param('cz_organization_id', 'org-2', 'org-2', id='cz_organization_id_param'),
        pytest.param('organizationId', 'org-3', 'org-3', id='organizationId_param'),
        pytest.param(None, None, ORGANIZATION_NONE, id='no_org_param'),
    ],
)
def test_cached_function_uses_correct_organization_param(context, param_name, org_id, expected_sk_org_id):
    cache_name = 'test-cache'

    if param_name == 'organization_id':

        @cached(cache_name)
        def test_function(organization_id: str, param: str) -> Dict[str, Any]:
            return {'result': 'value'}

    elif param_name == 'cz_organization_id':

        @cached(cache_name)
        def test_function(cz_organization_id: str, param: str) -> Dict[str, Any]:
            return {'result': 'value'}

    elif param_name == 'organizationId':

        @cached(cache_name)
        def test_function(organizationId: str, param: str) -> Dict[str, Any]:
            return {'result': 'value'}

    else:

        @cached(cache_name)
        def test_function(param: str) -> Dict[str, Any]:
            return {'result': 'value'}

    kwargs = {param_name: org_id, 'param': 'test'} if param_name and org_id else {'param': 'test'}
    test_function(**kwargs)

    expected_sk = f'{ORGANIZATION_KEY}#{expected_sk_org_id}#{TOKEN_KEY}'
    token_key = {'PK': f'{KEYSPACE}#{cache_name}', 'SK': expected_sk}
    context.mock_table.get_item.assert_called_once_with(Key=token_key, ConsistentRead=True)


@pytest.mark.unit
@pytest.mark.parametrize(
    'cached_data,expected_coerce_success,expected_result',
    [
        pytest.param({'name': 'test', 'value': 42}, True, TestResult('test', 42), id='valid_coercible_data'),
        pytest.param('not-a-dict', False, TestResult('fresh', 99), id='invalid_non_coercible_data'),
    ],
)
def test_cached_function_coerces_return_type(context, cached_data, expected_coerce_success, expected_result):
    cached_result = _CachedResult(token='token-123', parameters=None, data=cached_data)

    context.mock_result_store.get.return_value = cached_result
    context.mock_table.get_item.return_value = {'Item': {'data': 'token-123'}}

    @cached('test-cache')
    def test_function(param: str) -> TestResult:
        return TestResult(name='fresh', value=99)

    result = test_function('test')

    assert result == expected_result

    expected_result_store_set_calls = 0 if expected_coerce_success else 1
    assert context.mock_result_store.set.call_count == expected_result_store_set_calls


@pytest.mark.unit
def test_cached_name_validation_should_reject_invalid_names(context):
    with pytest.raises(ValueError):

        @cached('invalid#name')
        def test_function():
            pass


@pytest.mark.unit
def test_cached_should_use_lru_cache_for_memory_caching(context):
    @cached('test-cache', memcache_maxsize=2)
    def test_function(param: str) -> Dict[str, Any]:
        return {'result': f'computed-{param}'}

    result1 = test_function('same')
    result2 = test_function('same')

    assert result1 == result2 == {'result': 'computed-same'}
    assert context.mock_result_store.get.call_count == 1

    result3 = test_function('different')

    assert context.mock_result_store.get.call_count == 2
    assert result3 == {'result': 'computed-different'}


@pytest.mark.unit
@pytest.mark.parametrize(
    'specified_org,expected_invalidate_orgs',
    [
        pytest.param('test-org', ['test-org'], id='invalidate_specific_org'),
        pytest.param(None, ['org1', 'org2'], id='invalidate_all_orgs'),
    ],
)
def test_invalidate_sets_new_tokens(context, specified_org: str, expected_invalidate_orgs: List[str]):
    if not specified_org:
        context.mock_dyn.query.return_value = [
            {'SK': f'{ORGANIZATION_KEY}#org1#{TOKEN_KEY}'},
            {'SK': f'{ORGANIZATION_KEY}#org2#{TOKEN_KEY}'},
        ]

    invalidate('test-cache', specified_org)

    assert context.mock_dyn.put_item.call_args_list == [
        call(
            context.mock_table,
            Item={'PK': 'fast_distributed_cache#test-cache', 'SK': f'organization#{org}#token', 'data': 'test-uuid'},
        )
        for org in expected_invalidate_orgs
    ]

    def expr_str(condition):
        builder = ConditionExpressionBuilder()
        return builder.build_expression(condition, is_key_condition=True)

    if not specified_org:
        query_condition = Key('PK').eq(f'{KEYSPACE}#test-cache') & Key('SK').begins_with(f'{ORGANIZATION_KEY}#')

        context.mock_dyn.query.assert_called_once_with(
            context.mock_table,
            dict(KeyConditionExpression=MATCH(lambda x: expr_str(x) == expr_str(query_condition)), ConsistentRead=True),
        )


@pytest.mark.unit
@pytest.mark.parametrize(
    'existing_token,expect_new_token',
    [
        pytest.param({'Item': {'data': 'old-token'}}, False, id='token_exists'),
        pytest.param({}, True, id='token_does_not_exist'),
    ],
)
def test_cached_function_creates_token_when_needed(context, existing_token, expect_new_token):
    context.mock_table.get_item.return_value = existing_token

    @cached('test-cache')
    def test_function(param: str) -> Dict[str, Any]:
        return {'result': 'computed'}

    test_function('test')

    expected_put_item_calls = (
        [
            call(
                context.mock_table,
                Item={
                    'PK': f'{KEYSPACE}#test-cache',
                    'SK': f'{ORGANIZATION_KEY}#{ORGANIZATION_NONE}#{TOKEN_KEY}',
                    'data': 'test-uuid',
                },
            )
        ]
        if expect_new_token
        else []
    )
    assert context.mock_dyn.put_item.call_args_list == expected_put_item_calls


@pytest.mark.unit
def test_cached_function_returns_expected_values(
    context,
    lrucache_spy,
    # current_token: str, cached_data, expected_result, expected_set_called
):
    dynamo_data = {}
    result_store = {}

    def key(item: dict):
        return f'PK={item["PK"]}:SK={item["SK"]}'

    def get_item(Key: dict, ConsistentRead: bool):
        return dynamo_data.get(key(Key)) or {}

    def put_item(_table: Any, Item: dict):
        dynamo_data[key(Item)] = {'Item': Item}

    def query(_table: Any, _cond: Any):
        return [x['Item'] for x in dynamo_data.values()]

    def result_store_get(conn: Any, *key: str):
        return result_store.get(tuple(key))

    def result_store_set(conn: Any, *key: str, data: _CachedResult):
        result_store[tuple(key)] = data._replace(data={**data.data, '__from_dynamo__': True})

    context.mock_table.get_item.side_effect = get_item
    context.mock_dyn.put_item.side_effect = put_item
    context.mock_dyn.query.side_effect = query
    context.mock_result_store.get.side_effect = result_store_get
    context.mock_result_store.set.side_effect = result_store_set
    context.mock_uuid4.side_effect = uuid4

    def test_function_inner(param1: str, param2: int) -> Dict[str, Any]:
        return {'param1': param1, 'param2': param2}

    test_function_spy = Mock(wraps=test_function_inner)

    @cached('test-cache')
    def test_function(param1: str, param2: int) -> Dict[str, Any]:
        return test_function_spy(param1, param2)

    class Source(Enum):
        dynamo = 'dynamo'
        function = 'function'

    def check_cache(param1: str, param2: int, expected_source: Source):
        result = test_function(param1, param2)
        from_dynamo = {'__from_dynamo__': True} if expected_source == Source.dynamo else {}
        expected_result = {'param1': param1, 'param2': param2, **from_dynamo}
        assert result == expected_result
        if expected_source == Source.function:
            test_function_spy.assert_called_once_with(param1, param2)
        else:
            test_function_spy.assert_not_called()
        test_function_spy.reset_mock()

        result = test_function(param1, param2)
        assert result == expected_result
        test_function_spy.assert_not_called()

    check_cache('test', 123, expected_source=Source.function)
    check_cache('another test', 123, expected_source=Source.function)

    lrucache_spy.spy_return.clear()

    check_cache('test', 123, expected_source=Source.dynamo)
    check_cache('another test', 123, expected_source=Source.dynamo)

    fast_distributed_cache.invalidate('test-cache')

    check_cache('test', 123, expected_source=Source.function)
    check_cache('another test', 123, expected_source=Source.function)


@pytest.mark.unit
@pytest.mark.parametrize(
    'override,expected_result,expect_result_store_set',
    [
        pytest.param(lambda f: None, 'original value', True, id='no_override'),
        pytest.param(lambda f: lambda: 'overridden', 'overridden', False, id='override_function_with_fixed_value'),
        pytest.param(lambda f: f, 'original value', False, id='disable_cache'),
    ],
)
def test_override_cached_function_can_be_patched_to_disable_cache(
    context, mocker, override, expected_result, expect_result_store_set
):
    ocf = mocker.patch('czc.fast_distributed_cache.__override_cached_function__')

    @cached('test-cache')
    def test_function() -> str:
        return 'original value'

    ocf.side_effect = override
    assert test_function() == expected_result
    expected_result_store_set_calls = 1 if expect_result_store_set else 0
    assert context.mock_result_store.set.call_count == expected_result_store_set_calls
