import os
import sys

import loguru

_logger = None


class NullLogger:
    def __getattr__(self, name):
        return lambda *args, **kwargs: None

    def bind(self, **kwargs):
        return self


def setup_logger(exp_dir: str | None = None):
    global _logger

    loguru.logger.remove()
    if int(os.getenv("RANK", 0)) > 0:
        _logger = NullLogger()
        return _logger

    _logger = loguru.logger
    options = dict(
        level="INFO",
        backtrace=False,
        diagnose=False,
        format="{time:YYYY-MM-DD HH:mm:ss} | {level} | {message}",
    )
    _logger.add(sys.stderr, **options)
    if exp_dir is not None:
        _logger.add(os.path.join(exp_dir, "train.log"), colorize=False, encoding="utf-8", **options)
    return _logger


def get_logger():
    return _logger if _logger is not None else setup_logger()


__all__ = ["setup_logger", "get_logger"]