from typing import Union
import logging


def configure_logging(
        logging_level: Union[int, str] = logging.INFO,
        logging_fmt: str = "%(levelname)s: %(name)s: %(funcName)s: %(message)s",
        logger_obj: Union[None, logging.Logger] = None,
) -> logging.Logger:
    """
    Setup logging on the root logger, because `transformers` calls `logger.info` upon import.

    Adapted from:
        https://stackoverflow.com/a/54366471/5825811
        Configures a simple console logger with the given level.
        A use-case is to change the formatting of the default handler of the root logger.

    Format variables:
        https://docs.python.org/3/library/logging.html#logrecord-attributes
    """
    logger_obj = logger_obj or logging.getLogger()  # either the given logger or the root logger
    # logger_obj.removeHandler(stanza.log_handler)
    logger_obj.handlers.clear()
    logger_obj.setLevel(logging_level)
    # If the logger has handlers, we configure the first one. Otherwise we add a handler and configure it
    if logger_obj.handlers:
        console = logger_obj.handlers[0]  # we assume the first handler is the one we want to configure
    else:
        console = logging.StreamHandler()
        logger_obj.addHandler(console)
    console.setFormatter(logging.Formatter(logging_fmt))
    console.setLevel(logging_level)
    # # Work around to update TensorFlow's absl.logging threshold which alters the
    # # default Python logging output behavior when present.
    # # see: https://github.com/abseil/abseil-py/issues/99
    # # and: https://github.com/tensorflow/tensorflow/issues/26691#issuecomment-500369493
    # try:
    #     import absl.logging
    # except ImportError:
    #     pass
    # else:
    #     absl.logging.set_verbosity("info")
    #     absl.logging.set_stderrthreshold("info")
    #     absl.logging._warn_preinit_stderr = False
    return logger_obj