import logging, sys

logger_initialized = {}


def get_logger(name, with_stream=True, log_file=None, log_level=logging.INFO):
    """Initialize and get a logger by name.
    If the logger has not been initialized, this method will initialize the logger by adding one or two handlers,
    otherwise the initialized logger will be directly returned. During initialization, a StreamHandler will always be
    added. If `log_file` is specified and the process rank is 0, a FileHandler will also be added.

    Args:
        name (str): Logger name.
        log_file (str | None): The log filename. If specified, a FileHandler will be added to the logger.
        log_level (int): The logger level. Note that only the process of rank 0 is affected, and other processes will
            set the level to "Error" thus be silent most of the time.
    Returns:
        logging.Logger: The expected logger.
    """
    logger = logging.getLogger(name)
    if name in logger_initialized:
        return logger
    # handle hierarchical names
    # e.g., logger "a" is initialized, then logger "a.b" will skip the initialization since it is a child of "a".
    for logger_name in logger_initialized:
        if name.startswith(logger_name):
            return logger

    logger.propagate = False
    handlers = []
    if with_stream:
        stream_handler = logging.StreamHandler()
        handlers.append(stream_handler)

    import torch.distributed as dist
    if dist.is_available() and dist.is_initialized():
        rank = dist.get_rank()
    else:
        rank = 0
    # only rank 0 will add a FileHandler
    if rank == 0:
        if log_file is not None:
            file_handler = logging.FileHandler(log_file, 'w')
            handlers.append(file_handler)
    formatter = logging.Formatter(f'%(name)s - %(levelname)s - %(asctime)s - %(message)s', datefmt='%Y-%m-%d %H:%M:%S')
    logger.handlers = []
    for handler in handlers:
        handler.setFormatter(formatter)
        handler.setLevel(log_level)
        logger.addHandler(handler)
    if rank == 0:
        logger.setLevel(log_level)
    else:
        logger.setLevel(logging.ERROR)
    logger_initialized[name] = True
    return logger


def get_root_logger(log_file=None, log_level=logging.INFO):
    logger = get_logger(name='mani_skill_learn', log_file=log_file, log_level=log_level)
    return logger


def print_log(msg, logger=None, level=logging.INFO):
    """Print a log message.
    Args:
        msg (str): The message to be logged.
        logger (logging.Logger | str | None): The logger to be used.
            Some special loggers are:
            - "silent": no message will be printed.
            - other str: the logger obtained with `get_root_logger(logger)`.
            - None: The `print()` method will be used to print log messages.
        level (int): Logging level. Only available when `logger` is a Logger object or "root".
    """
    if logger is None:
        print(msg)
    elif isinstance(logger, logging.Logger):
        logger.log(level, msg)
    elif logger == 'silent':
        pass
    elif isinstance(logger, str):
        _logger = get_logger(logger)
        _logger.log(level, msg)
    else:
        raise TypeError(f'logger should be either a logging.Logger object, str, "silent" or None, '
                        f'but got {type(logger)}')


def flush_print(*args):
    print(*args)
    sys.stdout.flush()
