from datetime import datetime, timezone
from typing import Any

import boto3
from botocore import config as botocore_config
from dynamodb_encryption_sdk.delegated_keys.jce import JceNameLocalDelegatedKey
from dynamodb_encryption_sdk.encrypted import CryptoConfig
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 import CryptographicMaterialsProvider
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 mypy_boto3_dynamodb.service_resource import Table

from product_staging import config


def get_dynamodb_table() -> Table:
    """Get a dynamoDB table resource."""
    config_instance = botocore_config.Config()
    client = boto3.resource("dynamodb", config=config_instance)
    if config.DYNAMODB_TABLE_NAME:
        return client.Table(config.DYNAMODB_TABLE_NAME)
    else:
        raise Exception("No dynamodb table name defined")


def get_local_cmp() -> WrappedCryptographicMaterialsProvider:
    """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: str) -> AwsKmsCryptographicMaterialsProvider:
    """Get KMS cryptographic materials provider."""
    return AwsKmsCryptographicMaterialsProvider(kms_key_id)


def get_cryptographic_materials_provider() -> CryptographicMaterialsProvider:
    """Set up cryptographic materials provider."""
    kms_key_id = config.KMS_KEY_ID
    if kms_key_id:
        return 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
        return get_local_cmp()
    else:
        raise Exception(
            "Running without a secure CMP outside of dev. Configure KMS_KEY_ID env var."
        )


class DynamoDBManager:
    """DynamoDB table manager."""

    def __init__(self, table: Table):
        """Create Manager."""
        self.table = table
        self.partition_key = table.key_schema[0].get("AttributeName")
        self.crypto_provider = get_cryptographic_materials_provider()
        assert self.crypto_provider

    def crypto_config(self, attributes: dict[str, Any]) -> CryptoConfig:
        """Get crypto configuration."""
        encryption_context = EncryptionContext(
            table_name=self.table.table_name,
            partition_key_name=self.partition_key,
            attributes=attributes,
        )
        attribute_actions = AttributeActions(
            attribute_actions={
                self.partition_key: CryptoAction.DO_NOTHING,
                "created_at": CryptoAction.DO_NOTHING,
            },
            default_action=CryptoAction.ENCRYPT_AND_SIGN,
        )
        return CryptoConfig(
            materials_provider=self.crypto_provider,
            encryption_context=encryption_context,
            attribute_actions=attribute_actions,
        )

    def put_item(self, partition_key: str, data: dict[str, Any]) -> None:
        """Put encrypted item into dynamo table."""
        created_at = datetime.now(timezone.utc).strftime("%d/%m/%Y, %H:%M:%S")
        encrypted_item = self.encrypt_item(partition_key, data, created_at)
        self.table.put_item(Item=encrypted_item)

    def encrypt_item(
        self, partition_key: str, data: dict[str, Any], created_at: str
    ) -> dict[str, dict[str, Any]]:
        """Merge fields and encrypt the item."""
        primary_key = {self.partition_key: partition_key}
        ddb_attributes = dict_to_ddb(primary_key)
        crypto_config = self.crypto_config(ddb_attributes)
        merged_item = {
            **data,
            **primary_key,
            "created_at": created_at,
        }
        return encrypt_python_item(merged_item, crypto_config)

    def get_item(self, partition_key: str) -> dict[str, Any] | None:
        """Retrieve data from dynamo table."""
        record = self.table.get_item(Key={self.partition_key: partition_key})
        item = record.get("Item")
        if item is None:
            return None

        ddb_attributes = dict_to_ddb(item)
        crypto_config = self.crypto_config(ddb_attributes)
        return decrypt_python_item(item, crypto_config)