# Copyright (c) 2020 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, timedelta
from dateutil.tz import tzutc
from unittest.mock import MagicMock

import pytest
from boto3.dynamodb.conditions import Attr
from freezegun import freeze_time
from unittest.mock import ANY

from czc import locking
from czt.unittest import Context, MockEx, MATCH


def _table():
    table = MagicMock()
    MockEx(table.get_item).if_called_with(Key=ANY).return_value({})

    def mock_update(**kwargs):
        return {'Attributes': {'data': kwargs['ExpressionAttributeValues'][':data']}}

    table.update_item.side_effect = mock_update
    return table


def _data(lease_id='lease_id', expiration='2020-11-03T21:39:35.761200'):
    return {'lease_id': lease_id, 'expiration': expiration}


def _set_existing_lease(table, key, data):
    MockEx(table.get_item).if_called_with(Key=key).return_value({'Item': {'data': data}})


def _check_update(table, key=ANY, data=ANY, prior_lease_id=None, set_ttl=False):
    data_attr = Attr('data.lease_id')
    condition = data_attr.eq(prior_lease_id) if prior_lease_id is not None else data_attr.not_exists()

    def match_attr_values(v):
        return (v[':data'] == data and
                (not set_ttl or (v[':ttl'] > datetime.fromisoformat(v[':data']['expiration']).replace(tzinfo=tzutc()).timestamp())))

    assert MockEx(table.update_item).was_called_once_with(
        Key=key,
        ConditionExpression=condition,
        UpdateExpression='SET #data=:data' if not set_ttl else 'SET #data=:data, #ttl=:ttl',
        ExpressionAttributeNames={
            '#data': 'data',
            **({'#ttl': 'ttl'} if set_ttl else {})
        },
        ExpressionAttributeValues=MATCH(match_attr_values),
        ReturnValues='ALL_NEW'
    )


@pytest.mark.parametrize('key,hold_time,now,expected_expiration', [
    ({'PK': 'foo'}, timedelta(minutes=15), datetime(2020, 11, 3, 21, 51), '2020-11-03T22:06:00'),
    ({'PK': 'a', 'SK': 'b'}, timedelta(seconds=45), datetime(2020, 11, 3, 15, 10, 10), '2020-11-03T15:10:55'),
])
def test_acquire_lease_should_set_id_and_expiration_for_key_if_available(key, hold_time, now, expected_expiration):
    table = _table()
    with freeze_time(now):
        lease_id = locking.acquire_lease(table, key, hold_time)['lease_id']
    assert lease_id is not None
    _check_update(table, key=key, data=_data(lease_id=lease_id, expiration=expected_expiration))


@pytest.mark.parametrize('key,current_expiration,now,expected_expiration', [
    ({'PK': 'a', 'SK': 'b'}, '2020-11-03T15:20:00', datetime(2020, 11, 3, 15, 20, 1), '2020-11-03T15:25:01'),
    ({'PK': 'xyz'}, '2020-11-01T01:15:00', datetime(2020, 11, 1, 2, 15), '2020-11-01T02:20:00'),
])
def test_acquire_lease_should_set_id_and_expiration_if_expired(key, current_expiration, now, expected_expiration):
    table = _table()
    _set_existing_lease(table, key, _data(lease_id='prior', expiration=current_expiration))
    with freeze_time(now):
        lease_id = locking.acquire_lease(table, key, timedelta(minutes=5))['lease_id']
    assert lease_id is not None
    _check_update(table, key=key, data=_data(lease_id=lease_id, expiration=expected_expiration), prior_lease_id='prior')


@pytest.mark.parametrize('key,current_expiration,now', [
    ({'PK': 'a', 'SK': 'b'}, '2020-11-03T15:20:01', datetime(2020, 11, 3, 15, 20, 0)),
    ({'PK': 'xyz'}, '2020-11-01T02:15:00', datetime(2020, 11, 1, 1, 15)),
])
def test_acquire_lease_should_have_no_effect_if_valid_lease_held(key, current_expiration, now):
    table = _table()
    _set_existing_lease(table, key, _data(expiration=current_expiration))
    with freeze_time(now):
        lease_id = locking.acquire_lease(table, key, timedelta(minutes=5))
    assert lease_id is None
    assert not table.update_item.called


@pytest.mark.parametrize('key,hold_time,now,expected_expiration', [
    ({'PK': 'foo'}, timedelta(minutes=15), datetime(2020, 11, 3, 21, 51), '2020-11-03T22:06:00'),
    ({'PK': 'a', 'SK': 'b'}, timedelta(seconds=45), datetime(2020, 11, 3, 15, 10, 10), '2020-11-03T15:10:55'),
])
def test_acquire_lease_should_set_ttl_if_available_and_specified(key, hold_time, now, expected_expiration):
    table = _table()
    with freeze_time(now):
        lease_id = locking.acquire_lease(table, key, hold_time, True)['lease_id']
    assert lease_id is not None
    _check_update(table, key=key, data=_data(lease_id=lease_id, expiration=expected_expiration), set_ttl=True)


@pytest.mark.parametrize('key,lease_id', [
    ({'PK': 'foo', 'SK': 'bar'}, 'some-lease'),
    ({'PK': 'a', 'SK': 'b'}, 'other-lease-id'),
])
def test_release_lease_with_should_set_data_to_empty_when_within_expiration(key, lease_id):
    table = _table()
    _set_existing_lease(table, key, _data(lease_id=lease_id))
    with freeze_time(datetime(2020, 11, 3, 21, 38)):
        locking.release_lease(table, key, lease_id, datetime(2020, 11, 3, 21, 39))
        _check_update(table, key=key, data={}, prior_lease_id=lease_id)


@pytest.fixture(scope='function')
def context_manager_context(mocker):
    with Context(mocker, locking, ['acquire_lease', 'release_lease']) as context:
        yield context


@pytest.mark.unit
@pytest.mark.parametrize('key,hold_time,acquire_lock_result,mock_acquire_lease_return', [
    ({'PK': 'foo', 'SK': 'bar'}, timedelta(minutes=10), 'some-lease-id', {'lease_id': 'some-lease-id', 'expiration': '2020-11-03T21:48:00'}),
    ({'PK': 'a', 'SK': 'b'}, timedelta(seconds=30), None, None),
])
def test_lease_context_manager_should_acquire_lock(context_manager_context, key, hold_time, acquire_lock_result, mock_acquire_lease_return):
    table = MagicMock()
    context_manager_context.mock_acquire_lease.return_value = mock_acquire_lease_return
    with freeze_time(datetime(2020, 11, 3, 21, 38)):
        with locking.Lease(table, key, hold_time) as lease_result:
            assert MockEx(context_manager_context.mock_acquire_lease).was_called_once_with(table, key, hold_time, False)
            assert lease_result == acquire_lock_result


@pytest.mark.unit
@pytest.mark.parametrize('set_ttl', [True, False])
def test_lease_context_manager_should_acquire_lock_with_ttl(context_manager_context, set_ttl):
    table = MagicMock()
    key = {'PK': 'foo', 'SK': 'bar'}
    hold_time = timedelta(minutes=10)
    context_manager_context.mock_acquire_lease.return_value = {'lease_id': 'some-lease-id', 'expiration': '2020-11-03T21:48:00'}
    with freeze_time(datetime(2020, 11, 3, 21, 38)):
        with locking.Lease(table, key, hold_time, set_ttl) as lease_result:
            assert MockEx(context_manager_context.mock_acquire_lease).was_called_once_with(table, key, hold_time, set_ttl)
            assert lease_result == 'some-lease-id'


@pytest.mark.unit
@pytest.mark.parametrize('key,acquire_lock_result,should_release_lock', [
    ({'PK': 'foo', 'SK': 'bar'}, {'lease_id': 'some-lease-id', 'expiration': '2020-11-03T21:48:00'}, True),
    ({'PK': 'a', 'SK': 'b'}, None, False),
])
def test_lease_context_manager_should_release_lock_if_acquired(context_manager_context, key, acquire_lock_result,
                                                               should_release_lock):
    table = MagicMock()
    context_manager_context.mock_acquire_lease.return_value = acquire_lock_result
    with freeze_time(datetime(2020, 11, 3, 21, 38)):
        with locking.Lease(table, key, timedelta(seconds=1)):
            assert not context_manager_context.mock_release_lease.called

        if should_release_lock:
            assert MockEx(context_manager_context.mock_release_lease).was_called_with(table, key, acquire_lock_result['lease_id'], datetime.fromisoformat(acquire_lock_result['expiration']))
        else:
            assert not context_manager_context.mock_release_lease.called


@pytest.mark.unit
@pytest.mark.parametrize('key,expiration,now', [
    ({'PK': 'a', 'SK': 'b'}, datetime(2020, 11, 3, 15, 20, 3), datetime(2020, 11, 3, 15, 20, 4)),
    ({'PK': 'xyz'}, datetime(2020, 11, 1, 2, 15), datetime(2020, 11, 1, 3, 15)),
])
def test_realease_lease_should_raise_exception_when_trying_to_release_after_expiration(key, expiration, now):
    table = MagicMock()
    _set_existing_lease(table, key, _data())
    with freeze_time(now):
        with pytest.raises(locking.LeaseExpired):
            locking.release_lease(table, key, 'lease_id', expiration)
