from datetime import timedelta
from functools import wraps
from time import perf_counter
from typing import Any, Callable, Dict, Optional, Union


__all__ = ["PerfCounter", "timeit"]


class PerfCounter:
    def __init__(self, name: str):
        self.name = name
        self._counter: Dict[str, Dict[str, Union[int, float]]] = {}

    def add(self, key: str, value: float):
        if key not in self._counter:
            self._counter[key] = {"events": 1, "time": value}
        else:
            self._counter[key]["events"] += 1
            self._counter[key]["time"] += value

    @property
    def report(self) -> Dict[str, Any]:
        report = {}
        for k, v in self._counter.items():
            report[k] = {
                "events": v["events"],
                "time": timedelta(seconds=v["time"]),
                "avg": timedelta(seconds=(v["time"] / v["events"])),
            }
        return report


class timeit:
    def __init__(self, counter: Optional[PerfCounter] = None, key: Optional[str] = None):
        self._counter = counter
        self._key = key
        self._start = None

    def __call__(self, func: Callable):
        @wraps(func)
        def wrapped(ref, *args, **kwargs):
            self._key = func.__name__
            self._counter = ref.perf_counter
            with self:
                return func(ref, *args, **kwargs)

        return wrapped

    def __enter__(self):
        self._start = perf_counter()

    def __exit__(self, exc_type, exc_val, exc_tb):
        self._counter.add(self._key, perf_counter() - self._start)