import aws_encryption_sdk as encsdk  # pip install aws-encryption-sdk
import boto3
from aws_cryptographic_material_providers.mpl import AwsCryptographicMaterialProviders
from aws_cryptographic_material_providers.mpl.config import MaterialProvidersConfig
from aws_cryptographic_material_providers.mpl.models import CreateAwsKmsKeyringInput
from aws_cryptographic_material_providers.mpl.references import IKeyring
from aws_cryptographic_material_providers.smithygenerated.aws_cryptography_materialproviders.models import (
    CreateAwsKmsDiscoveryKeyringInput,
)
from dotenv import find_dotenv
from pydantic_settings import BaseSettings, SettingsConfigDict

from common.logging import get_logger

logger = get_logger(__name__)


class CryptoConfig(BaseSettings):
    key_id: str

    model_config = SettingsConfigDict(
        env_file=find_dotenv(), env_file_encoding="utf-8", extra="ignore", frozen=True
    )


class Crypto:
    def __init__(self, config: CryptoConfig):
        self.key_id = config.key_id
        self.client = encsdk.EncryptionSDKClient()
        kms_client = boto3.client("kms", region_name="us-east-1")

        mat_prov: AwsCryptographicMaterialProviders = AwsCryptographicMaterialProviders(
            config=MaterialProvidersConfig()
        )

        keyring_input: CreateAwsKmsKeyringInput = CreateAwsKmsKeyringInput(
            kms_key_id=self.key_id,
            kms_client=kms_client,
        )
        kms_keyring: IKeyring = mat_prov.create_aws_kms_keyring(input=keyring_input)
        self.keyring = kms_keyring
        discovery_keyring_input: CreateAwsKmsDiscoveryKeyringInput = (
            CreateAwsKmsDiscoveryKeyringInput(
                kms_client=kms_client,
            )
        )

        discovery_keyring: IKeyring = mat_prov.create_aws_kms_discovery_keyring(
            input=discovery_keyring_input
        )
        self.discovery_keyring = discovery_keyring
        logger.info(f"Keyring created with key ID: {self.key_id}")

    def encrypt(self, plaintext: bytes) -> bytes:
        """Encrypt the plaintext using the keyring."""

        ciphertext, _ = self.client.encrypt(source=plaintext, keyring=self.keyring)
        return ciphertext

    def decrypt(self, ciphertext: bytes) -> bytes:
        """Decrypt the ciphertext using the keyring."""

        plaintext, _ = self.client.decrypt(
            source=ciphertext,
            keyring=self.discovery_keyring,
        )
        return plaintext
