"""Secure document manager.

The secure document manager abstracts away the handling of the CMP, dynamoDB
client, and cryptographic context.

A simple API is offered to store and retrieve data.

"""

from functools import lru_cache
from typing import Type
from uuid import uuid4

from abacus_common_logic.utils.dates import current_timestamp
from boto3.dynamodb.conditions import Key
from dynamodb_encryption_sdk.delegated_keys.jce import JceNameLocalDelegatedKey
from dynamodb_encryption_sdk.encrypted import CryptoConfig
from dynamodb_encryption_sdk.encrypted.client import EncryptedPaginator
from dynamodb_encryption_sdk.encrypted.item import (
    decrypt_python_item,
    encrypt_python_item,
)
from dynamodb_encryption_sdk.identifiers import CryptoAction
from dynamodb_encryption_sdk.material_providers.aws_kms import (
    AwsKmsCryptographicMaterialsProvider,
)
from dynamodb_encryption_sdk.material_providers.wrapped import (
    WrappedCryptographicMaterialsProvider,
)
from dynamodb_encryption_sdk.structures import AttributeActions, EncryptionContext
from dynamodb_encryption_sdk.transform import dict_to_ddb

from payee.connectors.secure_data.document import SecureDocument


def get_local_cmp():
    """Get local cryptographic materials provider."""
    wrapping_key = JceNameLocalDelegatedKey.generate('AES', 256)
    signing_key = JceNameLocalDelegatedKey.generate('HmacSHA512', 512)
    return WrappedCryptographicMaterialsProvider(
        signing_key=signing_key, wrapping_key=wrapping_key, unwrapping_key=wrapping_key
    )


def get_kms_cmp(kms_key_id):
    """Get KMS cryptographic materials provider."""
    return AwsKmsCryptographicMaterialsProvider(kms_key_id)


class SecureDocumentManager:
    """Secure document manager."""

    current_document_version = '0'
    delimiter = '#'
    date_format = '%Y-%m-%dT%H:%M:%S.%f%z'

    def __init__(self, table, cryptographic_materials_provider):
        """Create Manager."""
        self.table = table
        # TODO - handle mis-configured tables not having a sort key.
        self.partition_key = table.key_schema[0].get('AttributeName')
        self.sort_key = table.key_schema[1].get('AttributeName')
        self.crypto_provider = cryptographic_materials_provider
        assert self.crypto_provider

    def crypto_config(self, attributes, document_field_actions):
        """Get crypto configuration."""
        encryption_context = self.get_encryption_context(attributes)
        attribute_actions = self.get_attribute_actions(document_field_actions)
        return self.build_crypto_config(encryption_context, attribute_actions)

    def build_crypto_config(self, encryption_context, attribute_actions):
        """Return the populated CryptoConfig."""
        return CryptoConfig(
            materials_provider=self.crypto_provider,
            encryption_context=encryption_context,
            attribute_actions=attribute_actions,
        )

    def get_encryption_context(self, attributes=None):
        """Get populated encryption context for fetching and storing."""
        return EncryptionContext(
            table_name=self.table.table_name,
            partition_key_name=self.partition_key,
            sort_key_name=self.sort_key,
            attributes=attributes,
        )

    def get_meta_attribute_actions(self):
        """Get meta fields as a dict that defines them as unencrypted.

        Encrypted fields are not included because the default action is to
        encrypt.

        The partition and sort keys cannot be encrypted if we are to
        store and retrieve the data using them as the primary key.

        The created_at field is not sensitive data and can be used to order
        search results.

        """
        return {
            self.partition_key: CryptoAction.DO_NOTHING,
            self.sort_key: CryptoAction.DO_NOTHING,
            'created_at': CryptoAction.DO_NOTHING,
        }

    def get_attribute_actions(self, document_field_actions=None):
        """Declare the encryption status for fields."""
        attribute_actions = self.get_meta_attribute_actions()
        if document_field_actions:
            attribute_actions = {**document_field_actions, **attribute_actions}
        return AttributeActions(
            attribute_actions=attribute_actions,
            default_action=CryptoAction.ENCRYPT_AND_SIGN,
        )

    def put_item(self, encrypted_item):
        """Put encrypted data into dynamo table."""
        self.table.put_item(Item=encrypted_item)

    def save_item(self, owner_type, owner_id, document):
        """Save a document."""
        versioned_key, current_key = self.generate_versioned_document_primary_keys(
            owner_type, owner_id, document.document_type()
        )
        created_at = self._get_current_timestamp()
        versioned_item = self.encrypt_item(versioned_key, document, created_at)
        current_item = self.encrypt_item(current_key, document, created_at)
        # TODO - use batch writer
        self.put_item(versioned_item)
        self.put_item(current_item)

    def purge_documents(self, owner_type, owner_id):
        """Purge all documents for a given entity."""

        with self.table.batch_writer() as batch:
            for item in self.table.query(
                KeyConditionExpression=Key(self.partition_key).eq(
                    self.build_partition_key(owner_type, owner_id)
                )
            )['Items']:
                batch.delete_item(
                    Key={
                        self.partition_key: item.get(self.partition_key),
                        self.sort_key: item.get(self.sort_key),
                    }
                )

    def save_versioned_item(
        self, owner_type: str, owner_id: int, document: SecureDocument
    ) -> str:
        """Save a new version without updating current."""
        version = uuid4()
        versioned_key = self.format_primary_key(
            owner_type, owner_id, document.document_type(), version
        )
        versioned_item = self.encrypt_item(
            versioned_key, document, self._get_current_timestamp()
        )
        self.put_item(versioned_item)
        return str(version)

    def _get_current_timestamp(self):
        """Get current timestamp."""
        return current_timestamp().strftime(self.date_format)

    def delete_current_item(
        self, owner_type: str, owner_id: int, document_class: Type[SecureDocument]
    ):
        """Delete the current document version."""
        current_key = self.format_primary_key(
            owner_type,
            owner_id,
            document_class.document_type(),
            self.current_document_version,
        )
        self.table.delete_item(Key=current_key)

    def generate_versioned_document_primary_keys(
        self, owner_type, owner_id, document_type
    ):
        """Generate the current and versioned primary keys."""
        versioned_key = self.format_primary_key(
            owner_type, owner_id, document_type, uuid4()
        )
        current_key = self.format_primary_key(
            owner_type, owner_id, document_type, self.current_document_version
        )
        return versioned_key, current_key

    def format_primary_key(self, owner_type, owner_id, document_type, version):
        """Return a formatted primary key."""
        return {
            self.partition_key: self.build_partition_key(owner_type, owner_id),
            self.sort_key: self.build_sort_key(document_type, version),
        }

    def extract_owner_partition_key(self, owner):
        """Return an owner_id of a given entity version from partition key.

        Extract owner_id from partition key eg.:
        `<owner_type><delimiter><owner_id>` -> [<owner_type>, <owner_id>]

        Args:
            owner (str): `owner` field (`PayeePayoneerDetailsDocument#221`)
        Return:
            (list): List of `owner_type` and `owner_id`
        """
        return owner.split(self.delimiter)

    def build_partition_key(self, owner_type, owner_id):
        """Format partition key."""
        return self.delimiter.join([owner_type, str(owner_id)])

    def build_sort_key(self, document_type, version=None):
        """Format sort key."""
        pieces = [str(i) for i in [document_type, version] if i]
        return self.delimiter.join(pieces)

    def encrypt_item(self, primary_key, document, created_at):
        """Merge fields and encrypt the item."""
        ddb_attributes = dict_to_ddb(primary_key)
        document_field_actions = document.get_field_actions()
        crypto_config = self.crypto_config(ddb_attributes, document_field_actions)
        merged_item = {
            **document.original_values,
            **primary_key,
            'created_at': created_at,
        }
        return encrypt_python_item(merged_item, crypto_config)

    def get_paginator(self, field_actions, method: str = 'query'):
        """Get the encrypted paginator for queries."""

        def crypto_config_method(**kwargs):
            """Return the crypto config for the pager."""
            return self.crypto_config(None, field_actions), kwargs

        return EncryptedPaginator(
            paginator=self.table.meta.client.get_paginator(method),
            decrypt_method=decrypt_python_item,
            crypto_config_method=crypto_config_method,
        )

    def get_items(self, owner_type, owner_id, document_class):
        """Get the current and previous versions of an item."""
        results = self._get_items(owner_type, owner_id, document_class)
        return self.extract_versions(results, document_class)

    def _get_items(self, owner_type, owner_id, document_class):
        """Get all the versions of an item."""
        params = self.get_search_key_condition(
            owner_type, owner_id, document_class().document_type()
        )
        paginator = self.get_paginator(document_class().get_field_actions())
        results = []
        for items in paginator.paginate(**params):
            for item in items.get('Items'):
                results.append(item)
        return results

    def get_item(self, owner_type, owner_id, document_class, revision=None):
        """Get the versioned item."""
        revision = revision if revision else self.current_document_version

        partition_key = self.build_partition_key(owner_type, owner_id)
        document = document_class()
        sort_key = self.build_sort_key(document.document_type(), revision)

        record = self.table.get_item(
            Key={self.partition_key: partition_key, self.sort_key: sort_key}
        )

        item = record.get('Item')
        if item is None:
            return None

        ddb_attributes = dict_to_ddb(item)
        field_actions = document.get_field_actions()
        crypto_config = self.crypto_config(ddb_attributes, field_actions)
        original_values = decrypt_python_item(item, crypto_config)
        document.set_values(original_values, set_empty_values=False)
        return document

    def get_search_key_condition(self, owner_type, owner_id, document_type):
        """Build the dynamo query search key."""
        key_condition = self.get_search_key(owner_type, owner_id, document_type)
        return {
            'TableName': self.table.table_name,
            'KeyConditionExpression': key_condition,
        }

    def extract_current_version(self, versions):
        """Get the current version from a list of versions."""
        current_item = None
        for item in versions:
            _, version = item.get(self.sort_key).split(self.delimiter)
            if version == self.current_document_version:
                current_item = item
        return current_item

    @staticmethod
    def extract_history(versions, current_item):
        """Get the history from a list of versions."""
        current_item_timestamp = current_item.get('created_at')
        history = []
        for item in versions:
            if current_item_timestamp != item.get('created_at'):
                history.append(item)
        return history

    @staticmethod
    def populate_closed_dates(versions):
        """Fill in the closed date on versions.

        Iterates through the list of versions and fills in the closed_at date
        on records using the created_at date of the next version.

        """
        last_created_at = None
        for item in reversed(sorted(versions, key=lambda i: i['created_at'])):
            if last_created_at:
                item['closed_at'] = last_created_at
            last_created_at = item.get('created_at')

    def extract_versions(self, versions, document_class):
        """Extract and format the versioned history of the document."""
        self.populate_closed_dates(versions)
        current_item = self.extract_current_version(versions)
        if not current_item:
            return None, []
        history = self.extract_history(versions, current_item)
        return document_class(current_item), self.history_from_list(
            history, document_class
        )

    @staticmethod
    def history_from_list(history, document_class):
        """Format a list of dicts as a list of documents."""
        documents = []
        for item in history:
            document = document_class(item)
            document.created_at = item.get('created_at')
            document.closed_at = item.get('closed_at')
            documents.append(document)
        return documents

    def get_search_key(self, owner_type, owner_id, document_type):
        """Build a dynamoDB table query.

        The query will search for exact matches on the partition key and
        a "begins with" match on the sort key.
        Returns the compound primary key.

        """
        partition_key = self.build_partition_key(owner_type, owner_id)
        sort_key = self.build_sort_key(document_type)
        p_key = Key(self.partition_key).eq(partition_key)
        s_key = Key(self.sort_key).begins_with(sort_key)
        return p_key & s_key

    def scan_items(self, document_class, scan_params) -> list:
        """Get the current versions of items by scan.

        Args:
            document_class (Class): class of items to scan
            scan_params (dict): params to scan prepared for dynamodb
                ex.: {
                    'FilterExpression': Attr('revision').eq(
                        'payee_payoneer_details#0')
                            & Attr('refresh_token_expires_at').lte(
                            '1654426075'),
                }
        Return:
            (list): found items by given criteria
        """
        scan_params['TableName'] = self.table.table_name
        paginator = self.get_paginator(document_class().get_field_actions(), 'scan')
        results = []
        for items in paginator.paginate(**scan_params):
            for item in items.get('Items'):
                results.append(item)
        return results


@lru_cache(maxsize=128)
def get_cryptographic_materials_provider(config):
    """Set up cryptographic materials provider."""
    from payee.utils.aws import access_kms_key

    cp_provider = False
    kms_key_id = config.SDM_CONFIG.get('KMS_KEY_ID')
    if kms_key_id:
        access_kms_key(kms_key_id)
        cp_provider = get_kms_cmp(kms_key_id)
    elif config.ENVIRONMENT not in [config.QA_ENVIRONMENT, config.PROD_ENVIRONMENT]:
        # only load a local cmp outside of deployed instances
        cp_provider = get_local_cmp()
    else:
        raise Exception(
            'Running without a secure CMP outside of dev. Configure KMS_KEY_ID env var.'
        )
    return cp_provider