"""Opt-in per-phase latency profile for :func:`deg.infer.fill`.

Off by default (zero cost: one ContextVar read per phase). Inside ``with Profile() as p``
every :func:`phase` synchronizes CUDA at its edges and accumulates wall time, call count
and the counters passed to :func:`count` (forward calls, rows, tokens).
"""

from __future__ import annotations

import time
from collections.abc import Iterator
from contextlib import contextmanager
from contextvars import ContextVar
from dataclasses import dataclass, field
from typing import Any

import torch

_ACTIVE: ContextVar[Profile | None] = ContextVar("deg_profile", default=None)


@dataclass
class Profile:
    seconds: dict[str, float] = field(default_factory=dict)
    calls: dict[str, int] = field(default_factory=dict)
    counters: dict[str, int] = field(default_factory=dict)
    _token: Any = None

    def __enter__(self) -> Profile:
        self._token = _ACTIVE.set(self)
        return self

    def __exit__(self, *exc: object) -> None:
        _ACTIVE.reset(self._token)

    def to_json(self) -> dict[str, Any]:
        return {"seconds": dict(self.seconds), "calls": dict(self.calls),
                "counters": dict(self.counters)}


def _sync() -> None:
    if torch.cuda.is_available() and torch.cuda.is_initialized():
        torch.cuda.synchronize()


@contextmanager
def phase(name: str) -> Iterator[None]:
    prof = _ACTIVE.get()
    if prof is None:
        yield
        return
    _sync()
    t0 = time.perf_counter()
    try:
        yield
    finally:
        _sync()
        prof.seconds[name] = prof.seconds.get(name, 0.0) + time.perf_counter() - t0
        prof.calls[name] = prof.calls.get(name, 0) + 1


def count(**values: int) -> None:
    prof = _ACTIVE.get()
    if prof is not None:
        for k, v in values.items():
            prof.counters[k] = prof.counters.get(k, 0) + int(v)