import aws_encryption_sdk as encsdk  # pip install aws-encryption-sdk
from aws_encryption_sdk.key_providers.kms import (
    DiscoveryAwsKmsMasterKeyProvider,
    StrictAwsKmsMasterKeyProvider,
)
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()
        self.encrypt_key_provider = StrictAwsKmsMasterKeyProvider(key_ids=[self.key_id])
        self.decrypt_key_provider = DiscoveryAwsKmsMasterKeyProvider()
        logger.info(f"KMS key providers initialized with key ID: {self.key_id}")

    def encrypt(self, plaintext: bytes) -> bytes:
        """Encrypt plaintext bytes with the configured KMS CMK."""

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

    def decrypt(self, ciphertext: bytes) -> bytes:
        """Decrypt ciphertext bytes with a KMS discovery provider."""

        plaintext, _ = self.client.decrypt(
            source=ciphertext,
            key_provider=self.decrypt_key_provider,
        )
        return plaintext
