# Copyright (c) 2021 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 functools import wraps

from aws_lambda_powertools import Logger
from aws_lambda_powertools.utilities.idempotency import DynamoDBPersistenceLayer, IdempotencyConfig, idempotent
from aws_lambda_powertools.utilities.idempotency.exceptions import IdempotencyItemNotFoundError
from aws_lambda_powertools.utilities.idempotency.persistence.base import DataRecord
from botocore.exceptions import ClientError
from jmespath import functions
from retrying import retry

from czc.api.v2.constants import CLOUDZERO_IDEMPOTENCY_KEY_HEADER, S3_INFO
from czc.dynamodb import get_table_name, is_max_item_size_exception
from czc.internal_api.client import RemoteApiError
from czc.s3 import get_bucket_name, get_json_object, put_json_object

logger = Logger()


# JMESPath does not support case insensitive key access;
# fortunately, we can add our own custom function to do that!
class CustomFunctions(functions.Functions):
    @functions.signature({'types': ['object']})
    def _func_get_cloudzero_idempotency_key_value(self, headers):
        name_lower = CLOUDZERO_IDEMPOTENCY_KEY_HEADER
        return next(
            # Iterate over the dict and do a case insensitive key comparison
            (value for key, value in headers.items() if key.lower() == name_lower),
            # None is returned if no matches were found
            None,
        )


class LargeResponseDynamoDBPersistenceLayer(DynamoDBPersistenceLayer):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._client = self._boto3_session.client('dynamodb', config=self._boto_config)

    def _save_response_data_to_s3(self, data_record: DataRecord):
        bucket_name = get_bucket_name()
        record_hash = self._generate_hash(data=data_record.idempotency_key)
        s3_key = f'{self.static_pk_value}/{record_hash}.json'
        put_json_object(bucket_name, s3_key, {self.data_attr: data_record.response_data})
        return {'M': {'key': {'S': s3_key}}}

    def _get_item_kwargs(self, data_record: DataRecord, s3_info=None):
        update_expression = 'SET #expiry = :expiry, #status = :status'
        expression_attr_values = {
            ':expiry': {'N': str(data_record.expiry_timestamp)},
            ':status': {'S': data_record.status},
        }

        expression_attr_names = {
            '#expiry': self.expiry_attr,
            '#status': self.status_attr,
        }

        if self.payload_validation_enabled:
            update_expression += ', #validation_key = :validation_key'
            expression_attr_values[':validation_key'] = {'S': data_record.payload_hash}  # type: ignore
            expression_attr_names['#validation_key'] = self.validation_key_attr

        if s3_info:
            update_expression += ', #s3_info = :s3_info'
            expression_attr_values[':s3_info'] = s3_info
            expression_attr_names['#s3_info'] = S3_INFO
        else:
            update_expression += ', #response_data = :response_data'
            expression_attr_values[':response_data'] = {'S': data_record.response_data}  # type: ignore
            expression_attr_names['#response_data'] = self.data_attr

        item_kwargs = {
            'Key': self._get_key(data_record.idempotency_key),
            'UpdateExpression': update_expression,
            'ExpressionAttributeValues': expression_attr_values,
            'ExpressionAttributeNames': expression_attr_names,
        }
        return item_kwargs

    def _update_record(self, data_record: DataRecord):
        logger.debug(f'Updating record for idempotency key: {data_record.idempotency_key}')
        item_kwargs = self._get_item_kwargs(data_record)

        try:
            self._client.update_item(TableName=self.table_name, **item_kwargs)  # type: ignore
        except ClientError as err:
            if is_max_item_size_exception(err):
                s3_info = self._save_response_data_to_s3(data_record)
                retry_item_kwargs = self._get_item_kwargs(data_record, s3_info=s3_info)
                self._client.update_item(TableName=self.table_name, **retry_item_kwargs)  # type: ignore
            else:
                raise

    def _get_record(self, idempotency_key) -> DataRecord:
        response = self._client.get_item(  # type: ignore
            TableName=self.table_name, Key=self._get_key(idempotency_key), ConsistentRead=True
        )

        try:
            item = response['Item']
        except KeyError as exc:
            raise IdempotencyItemNotFoundError from exc

        if S3_INFO in item:
            s3_info = self._deserializer.deserialize(item.pop(S3_INFO))
            s3_result = get_json_object(get_bucket_name(), s3_info['key'])
            item[self.data_attr] = {'S': s3_result[self.data_attr]}
        return self._item_to_data_record(item)


_default_config = IdempotencyConfig(
    event_key_jmespath='{ cloudzero_idempotency_key: not_null(headers | get_cloudzero_idempotency_key_value(@), requestContext.requestId), path: path, method: httpMethod }',
    jmespath_options={'custom_functions': CustomFunctions()},
)


def _get_default_persistence_store():
    return LargeResponseDynamoDBPersistenceLayer(
        expiry_attr='ttl',
        table_name=get_table_name(),
        key_attr='PK',
        static_pk_value=CLOUDZERO_IDEMPOTENCY_KEY_HEADER,
        sort_key_attr='SK',
    )


# NOTE: this does not use lambda_handler_decorator b/c idempotent() already returns a lambda_handler_decorator decorated middleware
def idempotency_key_aware(persistence_store: DynamoDBPersistenceLayer = None, **kwargs):
    """Make this Lambda Handler Idempotency Key Aware

    Parameters
    ----------
    persistence_store: DynamoDBPersistenceLayer, optional
        Custom DynamoDBPersistenceLayer, by default default_persistent_store

    Returns
    -------
    Callable
        Decorated Lambda Handler
    """
    persistence_store = persistence_store or _get_default_persistence_store()
    return idempotent(config=_default_config, persistence_store=persistence_store, **kwargs)  # type:ignore


def _on_idempotency_error(e):
    should_retry = isinstance(e, RemoteApiError) and 'IdempotencyAlreadyInProgressError' in str(e)
    logger.warning('Checking exception for retriability', extra={'should_retry': should_retry, 'exception': e})
    return should_retry


def retry_idempotency_error(wait_exponential_multiplier: int = 30, stop_max_attempt_number: int = 4):
    def retry_idempotency_error_handler(handler):
        @retry(
            retry_on_exception=_on_idempotency_error,
            wait_exponential_multiplier=wait_exponential_multiplier,
            stop_max_attempt_number=stop_max_attempt_number,
        )
        @wraps(handler)
        def wrapper(*args, **kwargs):
            return handler(*args, **kwargs)

        return wrapper

    return retry_idempotency_error_handler
