# 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 aws_lambda_powertools import Logger
from aws_lambda_powertools.utilities.idempotency import (IdempotencyConfig, idempotent, DynamoDBPersistenceLayer)
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 czc.api.v2.constants import CLOUDZERO_IDEMPOTENCY_KEY_HEADER, S3_INFO
from czc.dynamodb import get_table_name
from czc.s3 import get_bucket_name, put_json_object, get_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,
        )


def is_max_item_size_exception(e: ClientError) -> bool:
    error = e.response['Error']
    return (error['Code'] == 'ValidationException' and
            (error.get('Message', '').endswith('exceeded the maximum allowed size') or error.get('Message', '').endswith('exceeds max size'))
            )


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
