import atexit
import contextlib
import time
from collections.abc import Callable, Iterator
from functools import wraps
from statistics import mean, stdev
from timeit import timeit
from typing import Any, ContextManager, overload


def average_time(runnable: Callable[[], Any], runs: int) -> str:
    return analyze_results([_try_run(runnable) for _i in range(runs)])


def analyze_results(results: list[float | None]) -> str:
    times = [result for result in results if isinstance(result, float)]
    if len(times) == len(results):
        if len(times) > 2:
            deviation = f"{stdev(times):.2f}"
        else:
            deviation = "?"
        return f"success ${mean(times):.2f} \\pm {deviation}$"
    elif len(times) == 0:
        return "timeout/failure"
    else:
        return f"mixed ${mean(times):.2f}$ ({len(times)}/{len(results)} successes)"


def _try_run(runnable: Callable[[], Any]) -> float | None:
    try:
        return timeit(runnable, number=1)
    except:
        return None


_registered = False
_log: list[tuple[str, float]] = []


def _write_log() -> None:
    global _log
    for entry in _log:
        print(entry)


@overload
def log_perf[T, **P](
    f: Callable[P, T], *, threshold: float = 0.1
) -> Callable[P, T]: ...
@overload
def log_perf(f: str, *, threshold: float = 0.1) -> ContextManager[None]: ...


def log_perf[T, **P](
    f: str | Callable[P, T], *, threshold: float = 0.1
) -> Callable[P, T] | ContextManager[None]:
    global _log, _registered

    if not _registered:
        _registered = True
        atexit.register(_write_log)

    if callable(f):

        @wraps(f)
        def f_wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
            start = time.perf_counter()
            result = f(*args, **kwargs)
            end = time.perf_counter()
            diff = end - start

            if diff > threshold:
                _log.append((f.__name__, diff))
            return result

        return f_wrapper
    elif isinstance(f, str):

        @contextlib.contextmanager
        def yield_wrapper() -> Iterator[None]:
            start = time.perf_counter()
            yield
            end = time.perf_counter()
            diff = end - start

            if diff > threshold:
                _log.append((f, diff))

        return yield_wrapper()
