P = ParamSpec('P')
T = TypeVar('T')


@dataclass
class TimeContext:
    name: str
    start_ns: int
    end_ns: int | None = None

    _: KW_ONLY

    at_least_ns: int | None = None
    _generator: bool = False

    @property
    def duration_ns(self) -> int:
        if self.end_ns is None:
            return timer() - self.start_ns
        return self.end_ns - self.start_ns

    @property
    def should_report(self) -> bool:
        return self.at_least_ns is None or self.duration_ns >= self.at_least_ns

    def close(self) -> None:
        if self._generator:
            self._generator = False
            return
        self.end_ns = timer()
        if SHOW_TIMINGS and self.should_report:
            show_timing(self.name, self.start_ns, self.end_ns)

    def wrap_generator(self, gen: Iterator[T]) -> Iterator[T]:
        self._generator = True
        try:
            yield from gen
        finally:
            self.close()


@contextmanager
def timed_block(name: str, enable: bool = True) -> Iterator[TimeContext]:
    ctx = TimeContext(
        name, timer(), at_least_ns=100_000_000 if not enable else None
    )
    try:
        yield ctx
    finally:
        ctx.close()


def name_of(fn: Callable[..., Any]) -> str:
    if isinstance(fn, functools.partial):
        fn = fn.func
    return getattr(fn, '__qualname__', repr(fn))


def timed(fn: Callable[P, T]) -> Callable[P, T]:
    @functools.wraps(fn)
    def _wrapper(*args: P.args, **kw: P.kwargs) -> T:
        with timed_block(fn.__qualname__) as ctx:
            retval = fn(*args, **kw)
            if isinstance(retval, GeneratorType):
                return ctx.wrap_generator(retval)  # type: ignore
            else:
                return retval
    return _wrapper


def timed_generator(fn: Callable[P, Iterable[T]]) -> Callable[P, Iterable[T]]:
    @functools.wraps(fn)
    def _wrapper(*args: P.args, **kw: P.kwargs) -> Iterable[T]:
        with timed_block(fn.__qualname__):
            yield from fn(*args, **kw)
    return _wrapper