import time
from contextlib import contextmanager

import torch


class TemporalMemoryProfiler:
    def __init__(self, enabled=False, prefix='temporal_memory'):
        self.enabled = enabled
        self.prefix = prefix
        self.records = {}
        self.gpu_memory_allocated_mb = 0.0
        self.gpu_memory_max_allocated_mb = 0.0

    @contextmanager
    def measure(self, name):
        if not self.enabled:
            yield
            return
        self._sync()
        start = time.perf_counter()
        try:
            yield
        finally:
            self._sync()
            elapsed_ms = (time.perf_counter() - start) * 1000.0
            self.records[name] = self.records.get(name, 0.0) + elapsed_ms
            self._capture_memory()

    def report(self):
        if not self.enabled:
            return {}
        report = {
            f'{name}_latency_ms': round(value, 3)
            for name, value in self.records.items()
        }
        report['gpu_memory_allocated_mb'] = round(self.gpu_memory_allocated_mb, 3)
        report['gpu_memory_max_allocated_mb'] = round(self.gpu_memory_max_allocated_mb, 3)
        return report

    def log(self):
        report = self.report()
        if not report:
            return
        fields = ' '.join(f'{key}={value:.3f}' for key, value in report.items())
        print(f'[{self.prefix}] {fields}')

    def reset(self):
        self.records.clear()
        self.gpu_memory_allocated_mb = 0.0
        self.gpu_memory_max_allocated_mb = 0.0

    def _sync(self):
        if torch.cuda.is_available():
            torch.cuda.synchronize()

    def _capture_memory(self):
        if not torch.cuda.is_available():
            return
        mb = 1024.0 * 1024.0
        self.gpu_memory_allocated_mb = torch.cuda.memory_allocated() / mb
        self.gpu_memory_max_allocated_mb = torch.cuda.max_memory_allocated() / mb