# 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 uuid
from datetime import datetime, timedelta
from typing import Any, Callable, Optional

import pyfaaster.aws.tools as faaster_aws_tools
from boto3.dynamodb.conditions import Attr
from botocore.exceptions import ClientError
from retrying import retry
from toolz import get_in
from voluptuous import Schema

from czc.constants import PROJECT_NAME

logger = faaster_aws_tools.setup_logging(PROJECT_NAME)

LEASE_SCHEMA = Schema({
    'lease_id': str,
    'expiration': lambda x: datetime.fromisoformat(x)
})


def _is_conditional_check_exception(e):
    return isinstance(e, ClientError) and e.response['Error']['Code'] == 'ConditionalCheckFailedException'


@retry(retry_on_exception=_is_conditional_check_exception,
       stop_max_attempt_number=3,
       wait_random_min=1000,
       wait_random_max=2000)
def _update_lease(table: Any, key: dict, handler: Callable[[dict], Optional[dict]]) -> Optional[str]:
    existing_result = table.get_item(Key=key).get('Item')
    existing_data = LEASE_SCHEMA(existing_result['data']) if existing_result else {}
    updated_data = handler(existing_data)
    if updated_data is None:
        return None

    data_attr = Attr('data.lease_id')
    condition = data_attr.eq(existing_data['lease_id']) if 'lease_id' in existing_data else data_attr.not_exists()
    result = table.update_item(
        Key=key,
        ConditionExpression=condition,
        UpdateExpression='SET #data=:data',
        ExpressionAttributeNames={'#data': 'data'},
        ExpressionAttributeValues={':data': updated_data},
        ReturnValues='ALL_NEW'
    )
    return get_in(['Attributes', 'data', 'lease_id'], result)


def _lease_held(lease_data: dict) -> bool:
    return 'expiration' in lease_data and datetime.utcnow() < lease_data['expiration']


def acquire_lease(table: Any, key: dict, hold_time: timedelta) -> Optional[str]:
    """
    Attempts to acquire a time-bounded lock, stored under `key` in a DynamoDB `table`.
    Useful for standard critical-section type logic (e.g., start precisely one copy of
    some process) in a distributed context, where lock release can be unreliable.

    The lock becomes available again after (1) `release_lease` is called with the returned
    identifier, or (2) `hold_time` elapses. Due to condition (2), `hold_time` should be
    set to a value larger than the timeout of a Lambda where this is used.

    Args:
        table: a boto3 dynamodb `Table` resource
        key (dict): e.g., {'PK': 'organization_xyz', 'SK': 'lease_setup'}
        hold_time (timedelta): maximum time to hold the lease

    Returns:
        str or None: if the lock was successfully acquired, returns an identifer that
        can be used to release it; otherwise, returns None
    """

    def handler(lease_data):
        if not _lease_held(lease_data):
            expiration = (datetime.utcnow() + hold_time).isoformat()
            return {'lease_id': str(uuid.uuid4()), 'expiration': expiration}
        else:
            logger.info(f"Waiting on lease holder {lease_data['lease_id']} (expires {lease_data['expiration']})")
            return None

    return _update_lease(table, key, handler)


def release_lease(table: Any, key: dict, lease_id: str) -> None:
    """
    Release a lock acquired with `acquire_lease`, using the identifier returned from
    that function.

    Args:
        table: a boto3 dynamodb `Table` resource
        key (dict): e.g., {'PK': 'organization_xyz', 'SK': 'lease_setup'}
        lease_id (str): previously acquired identifier
    """

    def handler(lease_data):
        return {} if lease_data.get('lease_id') == lease_id else None

    _update_lease(table, key, handler)


class Lease:
    def __init__(self, table: Any, key: dict, hold_time: timedelta):
        self._table = table
        self._key = key
        self._hold_time = hold_time
        self._lease_id = None

    def __enter__(self):
        self._lease_id = acquire_lease(self._table, self._key, self._hold_time)
        if self._lease_id:
            logger.info(f'Acquired lease: {self._lease_id}')
        return self._lease_id

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self._lease_id is not None:
            release_lease(self._table, self._key, self._lease_id)
            logger.info(f'Released lease: {self._lease_id}')
