import logging
import time

logger = logging.getLogger("timing")


class TimeDiff:
    def __init__(self, name: str):
        self._name = name
        self._time_start = None
        self._silence = False

    async def __aenter__(self) -> "TimeDiff":
        return self.__enter__()

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        return self.__exit__(exc_type, exc_val, exc_tb)

    def __enter__(self) -> "TimeDiff":
        self._time_start = time.time()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        finish_time = time.time()
        if not self._silence:
            logger.info(
                "[%s] Total time: %.3f sec", self._name, finish_time - self._time_start
            )

    def muted(self):
        self._silence = True


def async_time_diff(func):
    async def wrapper(*args, **kwargs):
        async with TimeDiff(func.__qualname__):
            return await func(*args, **kwargs)

    return wrapper
