"""Generic checkpoint timer for instrumenting wall-clock/GPU-timeline durations."""
import time

import torch

from logging_utils import print_timing_divider, print_timing_lines


class Timer:
    """Records named marks and reports elapsed time between them.

    On CUDA, each mark also captures a CUDA event so elapsed_ms() reports true
    GPU-timeline duration instead of wall-clock time that includes async queuing.
    """

    def __init__(self, cuda_ok: bool) -> None:
        self._cuda_ok = cuda_ok
        self._marks: dict[str, tuple[torch.cuda.Event | None, float]] = {}

    def mark(self, name: str) -> None:
        ev: torch.cuda.Event | None = None
        if self._cuda_ok:
            ev = torch.cuda.Event(enable_timing=True)  # type: ignore[no-untyped-call]
            ev.record()
        self._marks[name] = (ev, time.perf_counter())

    def elapsed_ms(self, a: str, b: str) -> float:
        ev_a, t_a = self._marks[a]
        ev_b, t_b = self._marks[b]
        if ev_a and ev_b:
            return float(ev_a.elapsed_time(ev_b))  # true GPU-timeline ms
        return (t_b - t_a) * 1000.0

    def wall_start(self, name: str) -> float:
        return self._marks[name][1]

    def __contains__(self, name: str) -> bool:
        return name in self._marks

    def print_report(self, rows: list[tuple[str, str, str]], total: tuple[str, str] | None = None) -> None:
        """Print elapsed_ms() for each (label, start_mark, end_mark) row that has both marks recorded.

        If `total` is given as (start_mark, end_mark) and both are recorded, also prints an
        "overhead" line (total minus the sum of the printed rows) and a grand total line.
        """
        if self._cuda_ok:
            try:
                torch.cuda.synchronize()
            except Exception:
                pass
        total_ms = 0.0
        lines = []
        for label, a, b in rows:
            if a in self._marks and b in self._marks:
                ms = self.elapsed_ms(a, b)
                total_ms += ms
                lines.append(f"[timing]   {label:<14} {ms:8.1f} ms")
        if total and total[0] in self._marks and total[1] in self._marks:
            overall_ms = self.elapsed_ms(*total)
            lines.append(f"[timing]   {'overhead':<14} {overall_ms - total_ms:8.1f} ms")
            lines.append(f"[timing]   {'── total ──':<14} {overall_ms:8.1f} ms")
        print_timing_divider()
        print_timing_lines(lines)
        print_timing_divider()