import base64

import boto3

from app.app_types import EncryptedToken, PlainToken


class KMSClient:
    def __init__(self, region_name: str, key_id: str, session: boto3.Session) -> None:
        self.client = session.client("kms", region_name=region_name)
        self.key_id = key_id

    def decrypt(
        self,
        encrypted_token: EncryptedToken,
        *,
        context: dict[str, str] | None = None,
    ) -> PlainToken:
        context = context or {}
        decrypt_response = self.client.decrypt(
            KeyId=self.key_id,
            CiphertextBlob=base64.b64decode(encrypted_token),
            EncryptionContext=context,
        )

        return PlainToken(decrypt_response["Plaintext"].decode())