import time
import sys
from contextlib import contextmanager
from dataclasses import dataclass, field
from typing import Any, Iterator


@dataclass
class _Phase:
 name:     str
 start_ts: float
 end_ts:   float = 0.0


@dataclass
class StartupProfiler:
 _phases:  list[_Phase] = field(default_factory=list)
 _enabled: bool = True

 def enable(self) -> None:
  self._enabled = True

 def disable(self) -> None:
  self._enabled = False

 @contextmanager
 def phase(self, name: str) -> Iterator[None]:
  if not self._enabled:
   yield
   return
  p = _Phase(name=name, start_ts=time.perf_counter())
  try:
   yield
  finally:
   p.end_ts = time.perf_counter()
   self._phases.append(p)
   try:
    sys.stderr.write(f"[startup] {name}: {(p.end_ts - p.start_ts) * 1000.0:.1f}ms\n")
    sys.stderr.flush()
   except Exception:
    pass

 def snapshot(self) -> dict[str, Any]:
  phases: dict[str, dict[str, float]] = {}
  total = 0.0
  for p in self._phases:
   duration_ms = max(0.0, (p.end_ts - p.start_ts) * 1000.0) if p.end_ts else 0.0
   phases[p.name] = {
    "start_ts":    p.start_ts,
    "duration_ms": round(duration_ms, 3),
   }
   total += duration_ms
  return {
   "phases":   phases,
   "total_ms": round(total, 3),
   "count":    len(phases),
  }

 def reset(self) -> None:
  self._phases.clear()


startup_profiler: StartupProfiler = StartupProfiler()