import os
import boto3

SNS_TOPIC_ARN = os.environ.get('SNS_TOPIC_ARN')


class Notifier():
    def __init__(self, client=None, sns_topic_arn: str = ""):
        assert SNS_TOPIC_ARN or sns_topic_arn, 'SNS_TOPIC_ARN ENV VAR MUST BE SET'
        self.client = client if client else boto3.Session().client('sns')
        self.sns_topic = sns_topic_arn if sns_topic_arn else SNS_TOPIC_ARN

    def notify(self, message: str, subject: str = "", attributes: dict = None):
        if not attributes:
            attributes = {}
        message_attributes = {
            key: {
                'DataType': 'String',
                'StringValue': value
            } for key, value in attributes.items()
        }
        self.client.publish(
            TopicArn=self.sns_topic,
            Message=message,
            Subject=subject,
            MessageAttributes=message_attributes
        )


# if __name__ == '__main__':
#     notifier = Notifier(boto3.client('sns'), 'arn:aws:sns:us-east-1:490126044637:jmTest')
#     notifier.notify('Test Message', 'Test Sub', {'TestAttr': 'test val'})
