# 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
import datetime
from typing import Optional, Dict

import boto3
from aws_lambda_powertools import Logger
from aws_lambda_powertools.utilities.idempotency import (IdempotencyConfig, idempotent, BasePersistenceLayer)
from aws_lambda_powertools.utilities.idempotency.exceptions import IdempotencyItemNotFoundError, IdempotencyItemAlreadyExistsError
from aws_lambda_powertools.utilities.idempotency.persistence.base import DataRecord
from botocore.config import Config
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 DynamoDBPersistenceLayer(BasePersistenceLayer):
    def __init__(
        self,
        table_name: str,
        key_attr: str = "PK",
        key_attr_value: str = "idempotence",
        range_attr: str = "SK",
        expiry_attr: str = "expiration",
        status_attr: str = "status",
        data_attr: str = "data",
        validation_key_attr: str = "validation",
        boto_config: Optional[Config] = None,
        boto3_session: Optional[boto3.session.Session] = None,
    ):
        """
        Initialize the DynamoDB client

        Parameters
        ----------
        table_name: str
            Name of the table to use for storing execution records
        key_attr: str, optional
            DynamoDB attribute name for hash key, by default "PK"
        key_attr_value: str, optional
            DynamoDB attribute value for hash key, by default "idempotence"
        range_attr: str, optional
            DynamoDB attribute name for range key, by default "SK"
        expiry_attr: str, optional
            DynamoDB attribute name for expiry timestamp, by default "expiration"
        status_attr: str, optional
            DynamoDB attribute name for status, by default "status"
        data_attr: str, optional
            DynamoDB attribute name for response data, by default "data"
        boto_config: botocore.config.Config, optional
            Botocore configuration to pass during client initialization
        boto3_session : boto3.session.Session, optional
            Boto3 session to use for AWS API communication

        Examples
        --------
        **Create a DynamoDB persistence layer with custom settings**

            >>> from aws_lambda_powertools.utilities.idempotency import (
            >>>    idempotent, DynamoDBPersistenceLayer
            >>> )
            >>>
            >>> persistence_store = DynamoDBPersistenceLayer(table_name="idempotency_store")
            >>>
            >>> @idempotent(persistence_store=persistence_store)
            >>> def handler(event, context):
            >>>     return {"StatusCode": 200}
        """

        boto_config = boto_config or Config()
        session = boto3_session or boto3.session.Session()
        self._ddb_resource = session.resource("dynamodb", config=boto_config)
        self.table_name = table_name
        self.table = self._ddb_resource.Table(self.table_name)
        self.key_attr = key_attr
        self.key_attr_value = key_attr_value
        self.range_attr = range_attr
        self.expiry_attr = expiry_attr
        self.status_attr = status_attr
        self.data_attr = data_attr
        self.validation_key_attr = validation_key_attr
        super(DynamoDBPersistenceLayer, self).__init__()

    def _item_to_data_record(self, item: dict) -> DataRecord:
        """
        Translate raw item records from DynamoDB to DataRecord

        Parameters
        ----------
        item: Dict[str, Union[str, int]]
            Item format from dynamodb response

        Returns
        -------
        DataRecord
            representation of item

        """
        return DataRecord(
            idempotency_key=item[self.key_attr],
            status=item[self.status_attr],
            expiry_timestamp=item[self.expiry_attr],
            response_data=item.get(self.data_attr),  # type: ignore
            payload_hash=item.get(self.validation_key_attr),  # type: ignore
        )

    def _get_key(self, idempotency_key) -> Dict:
        return {
            self.key_attr: self.key_attr_value,
            self.range_attr: idempotency_key,
        }

    def _get_record(self, idempotency_key) -> DataRecord:
        response = self.table.get_item(Key=self._get_key(idempotency_key), ConsistentRead=True)

        try:
            item = response["Item"]
        except KeyError:
            raise IdempotencyItemNotFoundError

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

    def _put_record(self, data_record: DataRecord) -> None:
        item = {
            **self._get_key(data_record.idempotency_key),
            self.expiry_attr: data_record.expiry_timestamp,
            self.status_attr: data_record.status,
        }

        if self.payload_validation_enabled:
            item[self.validation_key_attr] = data_record.payload_hash

        now = datetime.datetime.now()
        try:
            logger.debug(f"Putting record for idempotency key: {data_record.idempotency_key}")
            self.table.put_item(
                Item=item,
                ConditionExpression='attribute_not_exists(#key) OR #now < :now',
                ExpressionAttributeNames={"#key": self.range_attr, "#now": self.expiry_attr},
                ExpressionAttributeValues={":now": int(now.timestamp())}
            )
        except self._ddb_resource.meta.client.exceptions.ConditionalCheckFailedException:
            logger.debug(f"Failed to put record for already existing idempotency key: {data_record.idempotency_key}")
            raise IdempotencyItemAlreadyExistsError
        except Exception as e:
            logger.error('Failed to put record', exc_info=True)
            raise e

    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.key_attr_value}/{record_hash}.json'
        put_json_object(bucket_name, s3_key, {self.data_attr: data_record.response_data})
        return {'key': s3_key}

    def _get_item_kwargs(self, data_record: DataRecord, s3_info=None):
        update_expression = "SET #expiry = :expiry, #status = :status"
        expression_attr_values = {
            ":expiry": data_record.expiry_timestamp,
            ":status": 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"] = data_record.payload_hash
            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'] = data_record.response_data
            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.table.update_item(**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.table.update_item(**retry_item_kwargs)  # type: ignore
            else:
                raise

    def _delete_record(self, data_record: DataRecord) -> None:
        logger.debug(f"Deleting record for idempotency key: {data_record.idempotency_key}")
        self.table.delete_item(Key=self._get_key(data_record.idempotency_key))


_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 DynamoDBPersistenceLayer(
            expiry_attr='ttl',
            table_name=get_table_name(),
            key_attr='PK',
            key_attr_value=CLOUDZERO_IDEMPOTENCY_KEY_HEADER,
            range_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
