import time
import psutil
import gpustat
from dataclasses import dataclass
from typing import Callable, TypeVar, Any, Tuple, Optional

T = TypeVar('T')

@dataclass
class MemorySnapshot:
    """Snapshot of memory usage at a point in time.

    Attributes:
        ram_used: RAM usage in megabytes
        gpu_memory_used: GPU memory usage in megabytes, None if GPU not available
    """
    ram_used: float
    gpu_memory_used: Optional[float] = None

def get_memory_usage():
    """Get current memory usage for RAM and GPU."""
    ram = psutil.Process().memory_info().rss / (1024 * 1024)  # MB
    gpu = None
    try:
        gpus = gpustat.GPUStatCollection.new_query()
        if gpus:
            gpu = sum(gpu.memory_used for gpu in gpus)
    except Exception as e:
        # Don't print error, just return None for GPU
        pass
    return MemorySnapshot(ram_used=ram, gpu_memory_used=gpu)

def measure_time(func: Callable[..., T]) -> Callable[..., Tuple[T, float]]:
    def wrapper(*args: Any, **kwargs: Any) -> Tuple[T, float]:
        start = time.perf_counter()
        result = func(*args, **kwargs)
        end = time.perf_counter()
        return result, end - start
    return wrapper
