# 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

import asyncio
import uuid
from datetime import datetime, timedelta
from time import sleep
from unittest.mock import ANY, MagicMock

import pytest
from boto3.dynamodb.conditions import Attr
from dateutil.tz import tzutc
from freezegun import freeze_time

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


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, tzinfo=tzutc()), '2020-11-03T22:06:00+00:00'),
        (
            {'PK': 'a', 'SK': 'b'},
            timedelta(seconds=45),
            datetime(2020, 11, 3, 15, 10, 10, tzinfo=tzutc()),
            '2020-11-03T15:10:55+00:00',
        ),
    ],
)
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+00:00',
            datetime(2020, 11, 3, 15, 20, 1, tzinfo=tzutc()),
            '2020-11-03T15:25:01+00:00',
        ),
        ({'PK': 'xyz'}, '2020-11-01T01:15:00+00:00', datetime(2020, 11, 1, 2, 15, tzinfo=tzutc()), '2020-11-01T02:20:00+00: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+00:00', datetime(2020, 11, 3, 15, 20, 0, tzinfo=tzutc())),
        ({'PK': 'xyz'}, '2020-11-01T02:15:00+00:00', datetime(2020, 11, 1, 1, 15, tzinfo=tzutc())),
    ],
)
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, tzinfo=tzutc()), '2020-11-03T22:06:00+00:00'),
        (
            {'PK': 'a', 'SK': 'b'},
            timedelta(seconds=45),
            datetime(2020, 11, 3, 15, 10, 10, tzinfo=tzutc()),
            '2020-11-03T15:10:55+00:00',
        ),
    ],
)
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, tzinfo=tzutc())):
        locking.release_lease(table, key, lease_id, datetime(2020, 11, 3, 21, 39, tzinfo=tzutc()))
        _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+00:00',
    }
    with freeze_time(datetime(2020, 11, 3, 21, 38, tzinfo=tzutc())):
        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+00: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, tzinfo=tzutc())):
        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, tzinfo=tzutc()),
            datetime(2020, 11, 3, 15, 20, 4, tzinfo=tzutc()),
        ),
        ({'PK': 'xyz'}, datetime(2020, 11, 1, 2, 15, tzinfo=tzutc()), datetime(2020, 11, 1, 3, 15, tzinfo=tzutc())),
    ],
)
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)


@pytest.fixture(scope='function')
def lease(mocker):
    return {'PK': 'foo', 'SK': f'bar-{uuid.uuid4()}'}


@pytest.mark.unit
def test_try_to_acquire_held_non_blocking_lease(mock_table, lease):
    now = get_current_utc_datetime()
    with locking.Lease(mock_table, lease, timedelta(seconds=10)) as outer_lease:
        assert outer_lease is not None
        with locking.Lease(mock_table, lease, timedelta(seconds=10)) as inner_lease:
            # This should not block and the lease should not be acquired
            assert inner_lease is None
            wait_time = get_current_utc_datetime() - now
            assert wait_time < timedelta(seconds=1)


@pytest.mark.unit
def test_acquire_held_blocking_lease(mock_table, lease):
    now = get_current_utc_datetime()
    first_lease_acquired = False

    def _acquire_first_lease():
        nonlocal first_lease_acquired
        with locking.BlockingLease(
            mock_table,
            lease,
            timedelta(seconds=10),
            timeout=timedelta(seconds=10),
        ) as lease_id:
            wait_time = get_current_utc_datetime() - now
            first_lease_acquired = True
            if lease_id:
                sleep(2)
            return wait_time, lease_id

    def _acquire_second_lease():
        nonlocal first_lease_acquired
        while not first_lease_acquired:
            sleep(0.1)

        with locking.BlockingLease(
            mock_table,
            lease,
            timedelta(seconds=10),
            timeout=timedelta(seconds=10),
        ) as lease_id:
            wait_time = get_current_utc_datetime() - now
            return wait_time, lease_id

    loop = asyncio.new_event_loop()
    task1 = loop.run_in_executor(None, _acquire_first_lease)
    task2 = loop.run_in_executor(None, _acquire_second_lease)
    results = loop.run_until_complete(asyncio.gather(task1, task2))

    wait_time, lease_id = results[0]
    assert wait_time < timedelta(seconds=1)
    assert lease_id is not None

    wait_time, lease_id = results[1]
    assert wait_time >= timedelta(seconds=2)
    assert lease_id is not None


@pytest.mark.unit
def test_blocking_lease_timeout(mock_table, lease):
    with locking.BlockingLease(mock_table, lease, timedelta(seconds=10), timeout=timedelta(seconds=1)):
        with pytest.raises(locking.AcquireLeaseTimeout):
            with locking.BlockingLease(mock_table, lease, timedelta(seconds=10), timeout=timedelta(seconds=1)):
                pass
