from __future__ import annotations

import time
from contextlib import contextmanager

from loguru import logger


class PerfTimer:
    def __init__(self, name: str = "") -> None:
        self.name = name
        self._start = 0.0
        self._elapsed = 0.0

    def start(self) -> PerfTimer:
        self._start = time.perf_counter()
        return self

    def stop(self) -> float:
        self._elapsed = time.perf_counter() - self._start
        return self._elapsed

    @property
    def elapsed_ms(self) -> float:
        return self._elapsed * 1000


@contextmanager
def measure(name: str = ""):
    timer = PerfTimer(name).start()
    yield timer
    elapsed = timer.stop()
    if name:
        logger.debug("{} took {:.1f}ms", name, timer.elapsed_ms)