# 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 jmespath import functions
from aws_lambda_powertools.utilities.idempotency import DynamoDBPersistenceLayer, IdempotencyConfig, idempotent

from czc.api.v2.constants import CLOUDZERO_IDEMPOTENCY_KEY_HEADER
from czc.dynamodb import get_table_name


# 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,
        )


_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()},
)

default_persistence_store = DynamoDBPersistenceLayer(
    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
    ----------
    persisence_store: DynamoDBPersistenceLayer, optional
        Custom DynamoDBPersistenceLayer, by default default_persistent_store

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