import atexit
import contextlib
import functools
from collections.abc import Callable, Generator, Iterable, Iterator, Mapping, Sequence
from types import TracebackType
from typing import Self

from loguru import logger

from liblaf import grapes

from . import TimerRecords, get_time


class timer(Mapping[str, float], contextlib.AbstractContextManager):  # noqa: N801
    """A class for measuring and recording execution time using various counters.

    This class provides a context manager and decorator for timing code execution. It supports multiple counters, logging, and generating human-readable reports of the recorded timings.

    Example: Using `timer` as a context manager
        ```python
        with timer():
            # Code to be timed
            time.sleep(1)
        ```

    Example: Using `timer` as a decorator
        ```python
        @timer(report_at_exit=True)
        def example_function():
            # Code to be timed
            time.sleep(1)


        example_function()
        ```

    Example: Tracking an iterable with `timer`
        ```python
        items = [1, 2, 3]
        for item in timer().track(items):
            # Code to be timed for each item
            time.sleep(1)
        ```
    """

    counters: Sequence[str] = ["perf", "process"]
    depth: int = 0
    label: str | None = None
    record_log_level: int | str | None = "DEBUG"
    records: TimerRecords
    report_log_level: int | str | None = "INFO"
    _end: dict[str, float]
    _start: dict[str, float]

    def __init__(
        self,
        label: str | None = None,
        *,
        counters: Sequence[str] = ["perf", "process"],
        depth: int = 0,
        record_log_level: int | str | None = "DEBUG",
        report_log_level: int | str | None = "INFO",
        report_at_exit: bool = False,
    ) -> None:
        """Initializes the timer object.

        Args:
            label: The label to be used for the timer.
            counters: The list of counters to be used for timing.
            depth: The depth of the timer.
            record_log_level: The log level for recording the timer records.
            report_log_level: The log level for reporting the timer records.
            report_at_exit: If `True`, the timer will be added to the global `TIMERS` list.
        """
        self.counters = counters
        self.depth = depth
        self.label = label
        self.record_log_level = record_log_level
        self.records = TimerRecords(default_key=self.counters[0])
        self.report_log_level = report_log_level
        self._end = {}
        self._start = {}
        if report_at_exit:
            TIMERS.append(self)

    def __getitem__(self, key: str) -> float:
        """Retrieve the elapsed time for a given key.

        Args:
            key: The key for which to retrieve the elapsed time.

        Returns:
            The elapsed time corresponding to the given key.
        """
        return self._end[key] - self._start[key]

    def __iter__(self) -> Iterator[str]:
        """Returns an iterator over the counters.

        Yields:
            An iterator over the string representations of the counters.
        """
        return iter(self.counters)

    def __len__(self) -> int:
        """Return the number of counters.

        Returns:
            The number of counters in the collection.
        """
        return len(self.counters)

    def __enter__(self) -> Self:
        """Enter the runtime context related to this object.

        This method is called when execution flow enters the context of the `with` statement. It initializes the timer by setting the label if it is not already set, increments the depth counter, and starts the timer.

        Returns:
            The instance of the timer object.
        """
        if self.label is None:
            self.label = grapes.caller_location(2)
        self.depth += 1
        self.start()
        return self

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_value: BaseException | None,
        traceback: TracebackType | None,
    ) -> None:
        """Exit the runtime context related to this object.

        This method is called when the execution leaves the context of the `with` statement that this object is used in. It ensures that the timer is stopped and the depth counter is decremented.

        Args:
            exc_type: The exception type if an exception was raised, otherwise None.
            exc_value: The exception instance if an exception was raised, otherwise None.
            traceback: The traceback object if an exception was raised, otherwise None.
        """
        self.end()
        self.depth -= 1

    def __call__[**P, T](self, fn: Callable[P, T]) -> Callable[P, T]:
        """A decorator that wraps a function to measure its execution time.

        Args:
            fn: The function to be wrapped.

        Returns:
            The wrapped function with timing functionality.
        """
        if self.label is None:
            self.label = grapes.full_qual_name(fn)

        @functools.wraps(fn)
        def wrapped(*args: P.args, **kwargs: P.kwargs) -> T:
            self.depth += 1
            with self:
                ret: T = fn(*args, **kwargs)
            self.depth -= 1
            return ret

        return wrapped

    @property
    def count(self) -> int:
        """The number of records in the timer."""
        return self.records.count

    @property
    def elapsed(self) -> float:
        """The elapsed time for the first counter."""
        return self[self.counters[0]]

    def column(self, key: str | None = None) -> Sequence[float]:
        """Retrieve a sequence of float values from the records.

        Args:
            key: The key to retrieve the sequence of floats. If None, the default key is used.

        Returns:
            A sequence of float values corresponding to the given key.
        """
        return self.records.column(key)

    def end(self) -> None:
        """Ends the timing for all counters, records the end times, and logs the record.

        This method performs the following actions:
        1. Iterates over all counter names and records their end times using `get_time`.
        2. Appends the current instance to the records list.
        3. Increments the depth counter.
        4. Logs the current record.
        5. Decrements the depth counter.
        """
        for name in self.counters:
            self._end[name] = get_time(name)
        self.records.append(self)
        self.depth += 1
        self.log_record()
        self.depth -= 1

    def human_report(self) -> str:
        """Generates a human-readable report of the recorded timings.

        Returns:
            A string containing the human-readable report of the timings.
        """
        return self.records.human_report(self.label)

    def human_record(self) -> str:
        """Generate a human-readable string representation of the recorded durations.

        This method iterates over the items in the instance, converts each duration to a human-readable format using `grapes.human_duration`, and constructs a formatted string. If a label is present, it is included at the beginning of the string.

        Returns:
            A human-readable string representation of the recorded durations.
        """
        text: str = f"{self.label} > " if self.label else ""
        for k, v in self.items():
            human: str = grapes.human_duration(v)
            text += f"{k}: {human}, "
        text = text.removesuffix(", ")
        return text

    def log_record(self, depth: int = 1) -> None:
        """Logs a record if the record_log_level is set.

        Args:
            depth: The depth to be added to the current depth for logging.
        """
        if not self.record_log_level:
            return
        logger.opt(depth=self.depth + depth).log(
            self.record_log_level, self.human_record()
        )

    def log_report(self, depth: int = 1) -> None:
        """Logs a human-readable report at the specified log level.

        Args:
            depth: The stack depth to use for the log entry.
        """
        if not self.report_log_level:
            return
        logger.opt(depth=depth).log(self.report_log_level, self.human_report())

    def max(self, key: str | None = None) -> float:
        """Retrieve the maximum value of the specified counter.

        Args:
            key: The key of the counter for which to retrieve the maximum value.

        Returns:
            The maximum value of the specified counter.
        """
        return self.records.max(key)

    def mean(self, key: str | None = None) -> float:
        """Calculate the mean of the specified counter.

        Args:
            key: The key of the counter for which to calculate the mean.

        Returns:
            The mean value of the specified counter.
        """
        return self.records.mean(key)

    def median(self, key: str | None = None) -> float:
        """Calculate the median of the specified counter.

        Args:
            key: The key of the counter for which to calculate the median.

        Returns:
            The median value of the specified counter.
        """
        return self.records.median(key)

    def min(self, key: str | None = None) -> float:
        """Retrieve the minimum value of the specified counter.

        Args:
            key: The key of the counter for which to retrieve the minimum value.

        Returns:
            The minimum value of the specified counter.
        """
        return self.records.min(key)

    def start(self) -> None:
        """Starts the timer for each counter in the `counters` list.

        This method iterates over the `counters` list and initializes the corresponding start time for each counter using the `get_time` function.
        """
        for name in self.counters:
            self._start[name] = get_time(name)

    def stdev(self, key: str | None = None) -> float:
        """Calculate the standard deviation of the specified counter.

        Args:
            key: The key of the counter for which to calculate the standard deviation.

        Returns:
            The standard deviation of the specified counter.
        """
        return self.records.stdev(key)

    def track[T](
        self, iterable: Iterable[T], *, log_report: bool = True
    ) -> Generator[T]:
        """Tracks the execution time of an iterable's items.

        Args:
            iterable: The iterable to be tracked.
            log_report: If `True`, logs the report after tracking.

        Yields:
            The items of the input iterable, one by one.
        """
        self.depth += 1
        if self.label is None:
            self.label = grapes.caller_location(2)
        for item in iterable:
            with self:
                yield item
        if log_report:
            self.log_report(self.depth + 1)
        self.depth -= 1


TIMERS: list[timer] = []


def exit_hook() -> None:
    for timer in TIMERS:
        if timer.records.count <= 1:
            continue
        timer.log_report()


atexit.register(exit_hook)
