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

from boto3.dynamodb.conditions import Attr
from botocore.exceptions import ClientError
from dateutil.tz import tzutc
from retrying import retry
from toolz import get_in
from voluptuous import Schema

from czc.constants import DYNAMODB_TTL_OVERRIDE
from czc.date import get_current_utc_datetime
from czc.dynamodb import DEFAULT_TTL_KEY, dynamo_safe
from czc.environ import environ_aware
from czc.logging import get_logger

logger = get_logger()

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


class LeaseExpired(Exception):
    pass


class AcquireLeaseTimeout(Exception):
    pass


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


@environ_aware(optional=[DYNAMODB_TTL_OVERRIDE])
@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], tuple[Optional[dict], Optional[datetime]]],
    DYNAMODB_TTL_OVERRIDE: Optional[str] = None,
) -> Optional[dict]:
    existing_result = table.get_item(Key=key).get('Item')
    existing_data = LEASE_SCHEMA(existing_result['data']) if existing_result and 'data' in existing_result else {}
    updated_data, ttl = 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' if ttl is None else 'SET #data=:data, #ttl=:ttl',
        ExpressionAttributeNames={'#data': 'data', **({'#ttl': (DYNAMODB_TTL_OVERRIDE or DEFAULT_TTL_KEY)} if ttl else {})},
        ExpressionAttributeValues={':data': updated_data, **({':ttl': dynamo_safe(ttl.timestamp())} if ttl else {})},
        ReturnValues='ALL_NEW',
    )
    return get_in(['Attributes', 'data'], result)


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


def acquire_lease(table: Any, key: dict, hold_time: timedelta, set_ttl: Optional[bool] = False) -> Optional[dict[str, 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:
        dict or None: if the lock was successfully acquired, returns a dict containing
        lease_id and expiration that can be used to release it; otherwise, returns None
    """

    def handler(lease_data):
        if not _lease_held(lease_data):
            expiration = get_current_utc_datetime() + hold_time
            return (
                {'lease_id': str(uuid.uuid4()), 'expiration': expiration.isoformat()},
                (expiration.replace(tzinfo=tzutc()) + timedelta(minutes=1)) if set_ttl else None,
            )
        logger.info(
            f'Lease held: {lease_data["lease_id"]}',
            extra={'lease_id': lease_data['lease_id'], 'expiration': lease_data['expiration']},
        )
        return None, None

    return _update_lease(table, key, handler)


def release_lease(table: Any, key: dict, lease_id: str, expiration: datetime) -> 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

    Note: Dynamodb.get_item is eventually consistent.
    This means that we are not guaranteed to read in the lease we expect.
    But if we've acquired the lock, that means that we've written to the dynamo table
    and we should be able to do whatever we want within the expiration time. If we're
    outside of expiration, then we raise because we have no guarantees.
    """

    def handler(lease_data):
        if expiration < get_current_utc_datetime():
            raise LeaseExpired()
        return {}, None

    _update_lease(table, key, handler)


class Lease:
    """
    A lease is a time-bounded lock, stored under `key` in a DynamoDB `table`.

    ***NOTE***
     This is a non-blocking lease, meaning that it will return immediately even if the lease is not available.
    If the lease is already acquired, the return value willl be `None` and the caller is responsible for
    checking if the lease is available.

    If you need a blocking lease, use `BlockingLease`.

    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
        set_ttl (bool): if True, the lease record will have a ttl set. This should be set to True. This parameter is
                    only provided for backwards compatibility.

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

    def __init__(self, table: Any, key: dict, hold_time: timedelta, set_ttl: Optional[bool] = False):
        self._table = table
        self._key = key
        self._hold_time = hold_time
        self._lease_id: Optional[str] = None
        self._expiration: Optional[datetime] = None
        self._set_ttl = set_ttl

    def __enter__(self):
        lease_data = acquire_lease(self._table, self._key, self._hold_time, self._set_ttl) or {}
        self._lease_id = lease_data.get('lease_id')
        if self._lease_id:
            self._expiration = datetime.fromisoformat(lease_data['expiration'])
            logger.info(f'Acquired lease: {self._lease_id}', extra={'lease_id': self._lease_id, 'expiration': self._expiration})
        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, self._expiration)  # type: ignore[arg-type]
            logger.info(f'Released lease: {self._lease_id}', extra={'lease_id': self._lease_id})


class BlockingLease(Lease):
    """
    A blocking lease is a lease that will block until the lease is available.
    If the lease is not available after the timeout, an `AcquireLeaseTimeout` exception is raised.

    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
        timeout (timedelta): maximum time to wait for the lease to be available

    Returns:
        str: The lease identifier

    Raises:
        AcquireLeaseTimeout: If the lease is not available after the timeout
    """

    def __init__(
        self,
        table: Any,
        key: dict,
        hold_time: timedelta,
        timeout: timedelta,
    ):
        super().__init__(table, key, hold_time, True)
        self._timeout = timeout

    def __enter__(self):
        start_time = get_current_utc_datetime()

        while (result := super().__enter__()) is None:
            if get_current_utc_datetime() - start_time > self._timeout:
                raise AcquireLeaseTimeout()
            time.sleep(0.5)

        return result
