import logging


class CustomFormatter(logging.Formatter):
    grey = "\x1b[38;20m"
    yellow = "\x1b[33;20m"
    red = "\x1b[31;20m"
    bold_red = "\x1b[31;1m"
    reset = "\x1b[0m"
    format = '%(asctime)s %(levelname)s:%(message)s'

    FORMATS = {
        logging.DEBUG: grey + format + reset,
        logging.INFO: grey + format + reset,
        logging.WARNING: yellow + format + reset,
        logging.ERROR: red + format + reset,
        logging.CRITICAL: bold_red + format + reset
    }

    def format(self, record):
        log_fmt = self.FORMATS.get(record.levelno)
        formatter = logging.Formatter(fmt=log_fmt, datefmt='%d/%m/%Y %I:%M')
        return formatter.format(record)


def setup_logger(logger: logging.Logger = None, handler: logging.Handler = None) -> logging.Logger:
    """
    Set up the logger with the handler and formatter.
    :param logger: The logger to be set up. If None, the root logger is used.
    :param handler: The handler to be set up. If None, a StreamHandler is used with the CustomFormatter.
    """
    current_handler = handler
    if current_handler is None:
        current_handler = logging.StreamHandler()
        current_handler.setFormatter(CustomFormatter())
    current_logger = logger
    if current_logger is None:
        current_logger = logging.getLogger()
    current_logger.addHandler(current_handler)
    return current_logger
