import logging
from copy import copy

from pycolib.ansicolor import highlight, BOLD, WHITE, YELLOW, RED


class ColoredFormatter(logging.Formatter):

    def __init__(self, fmt, datefmt):
        logging.Formatter.__init__(self, fmt, datefmt)
        self._colors = {
            'message': {
                'default': (None, None),
                'exception': (BOLD + YELLOW, RED),
                },
            'levelname': {
                'default': (BOLD + WHITE, None),
                'debug': (None, None),
                },
            'time': {
                'default': (BOLD + WHITE, None),
                'debug': (None, None),
                },
            }


    def formatTime(self, record, datefmt=None):
        """Override, add color as specified."""
        formatted = logging.Formatter.formatTime(self, record, datefmt)
        fg, bg = self.get_color('time', record)
        if fg or bg:
            formatted = highlight(1, fg, bg) + formatted + highlight(0)
        return formatted


    def format(self, record):
        """Override, add color as specified."""
        # work on copy, do not alter original record, which is processed by other formatters too
        record = copy(record)

        ### message color
        fg, bg = self.get_color('message', record)
        if fg or bg:
            record.msg = highlight(1, fg, bg) + record.msg + highlight(0)

        ### levelname color
        fg, bg = self.get_color('levelname', record)
        if fg or bg:
            record.levelname = highlight(1, fg, bg) + record.levelname + highlight(0)

        return logging.Formatter.format(self, record)


    def set_color(self, fieldname, levelname, fg=None, bg=None):
        """Specify color for parts of log message.

        Args:
            fieldname: Which field to apply: 'message' | 'levelname' | 'time'
            levelname: Log level to apply, case insensitive: 'debug' | 'info' etc.
                Custom levels are supported too.
                Special name 'default' is used for unspecified levels.

        """
        self._colors[fieldname][levelname.lower()] = (fg, bg)


    def get_color(self, fieldname, record):
        """Get specified color for field levelname (from record)."""
        colormap = self._colors[fieldname]
        levelname = record.levelname.lower()
        fg, bg = colormap['default']
        if levelname in colormap:
            fg, bg = colormap[levelname]
        # exception is not real level, detect it using exc_info
        if record.levelno == logging.ERROR and record.exc_info is not None \
        and 'exception' in colormap:
            fg, bg = colormap['exception']
        return fg, bg