import os
import sys
import traceback
from logging import StreamHandler
from psi_notifier.notifier import Notifier


class PSINotifierLoggingHandler(StreamHandler):
    def __init__(self, notifier: Notifier = None, topic_arn: str = None):
        """Initialize a PSI Notifier logging handler
        If a Notifier object is supplied, this handler will use the provided
        Notifier.
        If a Notifier is not supplied, a topic_arn must be set for this handler
        to initialize it's own Notifier with the provided topic_arn

        Parameters
        ----------
        notifier: Notifier
            A pre-configured PSI Notifier object

        topic_arn: str
            The ARN to publish SNS messages to if notifier is not provided

        Example Usage
        -------------
        logger = logging.getLogger(os.path.basename(__file__))
        logger.setLevel(logging.INFO)
        formatter = logging.Formatter('<%(asctime)s> [%(levelname)s] %(name)s:%(funcName)s -> %(message)s')
        notifier = Notifier(sns_topic_arn='arn:aws:sns:us-east-1:490126044637:jmTest')
        notifier_handler = PSINotifierLoggingHandler(notifier)
        terminal_handler = logging.StreamHandler(sys.stdout)
        notifier_handler.setFormatter(formatter)
        notifier_handler.setLevel(logging.WARNING)
        terminal_handler.setFormatter(formatter)
        logger.addHandler(notifier_handler)
        logger.addHandler(terminal_handler)
        logger.info('This will only log to stdout')
        logger.warning('This will  log to SNS')
        """
        StreamHandler.__init__(self)
        if not notifier and not topic_arn:
            raise Exception('Must provide either a Notifier or SNS topic ARN!')
        if notifier and not topic_arn:
            self.notifier = notifier
        if topic_arn and not notifier:
            self.notifier = Notifier(sns_topic_arn=topic_arn)
        if topic_arn and notifier:
            self.notifier = Notifier(sns_topic_arn=topic_arn)

    def emit(self, record):
        message = self.format(record)
        self.notifier.notify(message, f'[{record.levelname}] {record.filename}:{record.funcName}:{record.lineno}')


def custom_exception_hook(notifier: Notifier):
    def handler(
            exception_type: type,
            exception_value: Exception,
            traceb,
            attributes: dict = None):
        except_file = traceback.TracebackException(exception_type,
                                                   exception_value,
                                                   traceb).stack[-1].filename
        notifier_params = {
            'subject': f'{os.path.basename(except_file)} Exception: {exception_type.__name__}: {exception_value}',
            'message': ''.join(traceback.format_tb(traceb)),
            'attributes': attributes
        }

        notifier.notify(**notifier_params)
        sys.__excepthook__(exception_type, exception_value, traceb)
    return handler
