# -*- coding: utf-8 -*-
# Copyright (c) 2017 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 logging
import uuid
from datetime import timedelta, datetime
from typing import Any, Callable, Optional

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 import constants

logger = logging.getLogger(constants.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:
            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)
        return self._lease_id

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self._lease_id is not None:
            return release_lease(self._table, self._key, self._lease_id)
