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."""

        ### 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)

    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