from __future__ import annotations

import base64
from typing import TYPE_CHECKING

import boto3
from fansifter_common.adapters.encrypter import EncryptedStr

if TYPE_CHECKING:
    from mypy_boto3_kms import KMSClient as _KMSClient


class KMSClient:
    def __init__(
        self, *, region_name: str, key_id: str, session: boto3.Session | None = None
    ):
        self.region_name = region_name
        self.key_id = key_id
        self._session = session
        self._client: _KMSClient | None = None

    @property
    def session(self) -> boto3.Session:
        if self._session is None:
            self._session = boto3.Session()
        return self._session

    @property
    def client(self) -> _KMSClient:
        if self._client is None:
            self._client = self.session.client("kms", region_name=self.region_name)
        return self._client

    def encrypt(self, value: str) -> EncryptedStr:
        response = self.client.encrypt(KeyId=self.key_id, Plaintext=value)
        decoded = base64.b64encode(response["CiphertextBlob"]).decode()
        return EncryptedStr(decoded)

    def decrypt(self, value: EncryptedStr) -> str:
        response = self.client.decrypt(
            KeyId=self.key_id,
            CiphertextBlob=base64.b64decode(value),
        )
        return response["Plaintext"].decode()