import functools
import logging
import os
import sys
from termcolor import colored


class _ColorfulFormatter(logging.Formatter):
    def __init__(self, *args, **kwargs):
        self._root_name = kwargs.pop('root_name') + '.'
        self._abbrev_name = kwargs.pop('abbrev_name', '')
        if len(self._abbrev_name):
            self._abbrev_name = self._abbrev_name + '.'
        super(_ColorfulFormatter, self).__init__(*args, **kwargs)

    def formatMessage(self, record):
        record.name = record.name.replace(self._root_name, self._abbrev_name)
        log = super(_ColorfulFormatter, self).formatMessage(record)
        if record.levelno == logging.WARNING:
            prefix = colored('WARNING', 'red', attrs=['blink'])
        elif record.levelno == logging.ERROR or record.levelno == logging.CRITICAL:
            prefix = colored('ERROR', 'red', attrs=['blink', 'underline'])
        else:
            return log
        return prefix + ' ' + log


# so that calling setup_logger multiple times won't add many handlers
@functools.lru_cache()
def setup_logger(output=None,
                 distributed_rank=0,
                 *,
                 color=True,
                 name='imagenet',
                 abbrev_name=None):
    """Initialize the detectron2 logger and set its verbosity level to "INFO".

    Args:
        output (str): a file name or a directory to save log. If None, will not save log file.
            If ends with ".txt" or ".log", assumed to be a file name.
            Otherwise, logs will be saved to `output/log.txt`.
        name (str): the root module name of this logger

    Returns:
        logging.Logger: a logger
    """
    logger = logging.getLogger(name)
    logger.setLevel(logging.DEBUG)
    logger.propagate = False

    if abbrev_name is None:
        abbrev_name = name

    plain_formatter = logging.Formatter(
        '[%(asctime)s.%(msecs)03d]: %(message)s', datefmt='%m/%d %H:%M:%S')
    # stdout logging: master only
    if distributed_rank == 0:
        ch = logging.StreamHandler(stream=sys.stdout)
        ch.setLevel(logging.DEBUG)
        if color:
            formatter = _ColorfulFormatter(
                colored('[%(asctime)s.%(msecs)03d]: ', 'green') +
                '%(message)s',
                datefmt='%m/%d %H:%M:%S',
                root_name=name,
                abbrev_name=str(abbrev_name),
            )
        else:
            formatter = plain_formatter
        ch.setFormatter(formatter)
        logger.addHandler(ch)

    # file logging: all workers
    if output is not None:
        if output.endswith('.txt') or output.endswith('.log'):
            filename = output
        else:
            filename = os.path.join(output, 'log.txt')
        if distributed_rank > 0:
            filename = filename + f'.rank{distributed_rank}'
        os.makedirs(os.path.dirname(filename), exist_ok=True)

        fh = logging.StreamHandler(_cached_log_stream(filename))
        fh.setLevel(logging.DEBUG)
        fh.setFormatter(plain_formatter)
        logger.addHandler(fh)

    return logger


@functools.lru_cache(maxsize=None)
def _cached_log_stream(filename):
    return open(filename, 'a')