import logging
import sys


class LsassyLogger(logging.LoggerAdapter):
    def __init__(self, no_color=False):
        super().__init__(self, extra=None)
        self.logger = logging.getLogger("lsassy")
        self.logger.propagate = False
        self.no_color = no_color

        formatter = LsassyFormatter(no_color=no_color)
        self.handler = logging.StreamHandler(sys.stdout)
        self.handler.setFormatter(formatter)
        self.logger.addHandler(self.handler)

    def lsassy_highlight(self, msg):
        """
        Highlight in yellow provided message
        :param msg: Message to lsassy_highlight
        :return: Highlighted message
        """
        if self.no_color:
            return msg
        return "\033[1;33m{}\033[0m".format(msg)

    def set_no_color(self, no_color=False):
        self.logger.removeHandler(self.handler)
        self.handler = logging.StreamHandler(sys.stdout)
        self.handler.setFormatter(LsassyFormatter(no_color=no_color))
        self.logger.addHandler(self.handler)


class LsassyFormatter(logging.Formatter):
    """
    Custom formatting. Inspired by impacket "Logger" class
    """
    def __init__(self, no_color=False):
        self.formatter = logging.Formatter.__init__(self, '%(bullet)s %(threadName)s %(message)s', None)
        self.no_color = no_color
        self.BLUE, self.WHITE, self.YELLOW, self.RED, self.GREEN, self.NC = '', '', '', '', '', ''
        if not self.no_color:
            self.BLUE = '\033[1;34m'
            self.WHITE = '\033[1;37m'
            self.YELLOW = '\033[1;33m'
            self.RED = '\033[1;31m'
            self.GREEN = '\033[1;32m'
            self.NC = '\033[0m'

    def format(self, record):
        """
        Custom bullet formatting with colors
        :param record: Record to log
        """
        if record.levelno == logging.INFO:
            record.bullet = '{}[*]{}'.format(self.BLUE, self.NC)
        elif record.levelno == logging.DEBUG:
            record.bullet = '{}[*]{}'.format(self.WHITE, self.NC)
        elif record.levelno == logging.WARNING:
            record.bullet = '{}[!]{}'.format(self.YELLOW, self.NC)
        elif record.levelno == logging.ERROR:
            record.bullet = '{}[x]{}'.format(self.RED, self.NC)
        else:
            record.bullet = '{}[+]{}'.format(self.GREEN, self.NC)

        # Only log stacktrace when log level is DEBUG
        if record.exc_info and logging.getLogger("lsassy").getEffectiveLevel() != logging.DEBUG:
            record.exc_info = None

        return logging.Formatter.format(self, record)


lsassy_logger = LsassyLogger()
