#!/usr/bin/env python3
"""
memory_snapshot.py - Unified Memory Snapshot Utility

WO-WOW-METRICS-01: Consolidated memory measurement for streaming/transport benchmarks.

Consolidates patterns from:
- stream_xw_matmul.py: get_rss_mb() via /proc/self/status
- bench_coldstart_real.py: read_proc_io() via /proc/self/io
- bench_modes.py: torch.cuda GPU memory tracking
- resource.getrusage(): page fault tracking
- /proc/self/smaps_rollup: working set (Pss) estimation

Usage:
    from helix_cdc.utils.memory_snapshot import take_snapshot, compute_delta

    baseline = take_snapshot()
    # ... do work ...
    final = take_snapshot()
    delta = compute_delta(baseline, final)
    print(f"RSS delta: {delta.rss_delta_mb:.1f} MB, Peak: {delta.rss_peak_mb:.1f} MB")
"""

from __future__ import annotations

import resource
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, Optional


@dataclass
class MemorySnapshot:
    """
    Point-in-time memory snapshot.

    Captures:
    - RSS (Resident Set Size) from /proc/self/status
    - GPU allocated/reserved memory from torch.cuda
    - I/O bytes read from /proc/self/io
    - Page faults from getrusage()
    - Working set estimate from /proc/self/smaps_rollup

    WO-TRUTHY-METRICS-01: Added io_rchar/io_wchar for cache-aware I/O metrics.
    """
    rss_mb: float = 0.0              # /proc/self/status VmRSS (kB -> MB)
    gpu_allocated_mb: float = 0.0    # torch.cuda.memory_allocated()
    gpu_reserved_mb: float = 0.0     # torch.cuda.memory_reserved()
    gpu_peak_mb: float = 0.0         # torch.cuda.max_memory_allocated()
    io_read_bytes: int = 0           # /proc/self/io read_bytes (actual disk I/O)
    io_write_bytes: int = 0          # /proc/self/io write_bytes
    io_rchar: int = 0                # /proc/self/io rchar (logical reads, includes cache)
    io_wchar: int = 0                # /proc/self/io wchar (logical writes)
    major_faults: int = 0            # ru_majflt from getrusage (disk reads)
    minor_faults: int = 0            # ru_minflt from getrusage (page reclaims)
    working_set_mb: float = 0.0      # /proc/self/smaps_rollup Pss
    timestamp: str = field(
        default_factory=lambda: datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")
    )

    def to_dict(self) -> Dict[str, Any]:
        return {
            "rss_mb": round(self.rss_mb, 2),
            "gpu_allocated_mb": round(self.gpu_allocated_mb, 2),
            "gpu_reserved_mb": round(self.gpu_reserved_mb, 2),
            "gpu_peak_mb": round(self.gpu_peak_mb, 2),
            "io_read_bytes": self.io_read_bytes,
            "io_write_bytes": self.io_write_bytes,
            "io_rchar": self.io_rchar,
            "io_wchar": self.io_wchar,
            "io_read_mb": round(self.io_read_bytes / (1024 * 1024), 2),
            "io_read_logical_mb": round(self.io_rchar / (1024 * 1024), 2),
            "major_faults": self.major_faults,
            "minor_faults": self.minor_faults,
            "working_set_mb": round(self.working_set_mb, 2),
            "timestamp": self.timestamp,
        }


@dataclass
class MemoryDelta:
    """
    Memory change between two snapshots.

    Useful for before/after measurement of operations.

    WO-TRUTHY-METRICS-01: Added bytes_read_disk, bytes_read_logical, cache_hit_ratio
    for cache-aware I/O metrics.
    """
    rss_delta_mb: float = 0.0        # RSS change (can be negative)
    rss_peak_mb: float = 0.0         # Peak RSS observed
    gpu_delta_mb: float = 0.0        # GPU allocated change
    gpu_peak_allocated_mb: float = 0.0  # Peak GPU allocated
    gpu_peak_reserved_mb: float = 0.0   # Peak GPU reserved
    bytes_read: int = 0              # I/O bytes read during interval (actual disk I/O)
    bytes_written: int = 0           # I/O bytes written during interval
    bytes_read_logical: int = 0      # Logical bytes read (includes cache hits)
    bytes_written_logical: int = 0   # Logical bytes written
    cache_hit_ratio: float = 0.0     # (logical - disk) / logical (0.0 to 1.0)
    page_faults: int = 0             # Total page faults (major + minor)
    major_faults: int = 0            # Disk I/O page faults
    minor_faults: int = 0            # In-memory page reclaims
    working_set_estimate_mb: float = 0.0  # Pss-based working set

    # Efficiency metrics (filled by compute_delta if model_size provided)
    model_size_mb: float = 0.0
    working_set_ratio: float = 0.0   # working_set / model_size
    read_vs_size_ratio: float = 0.0  # bytes_read / model_size

    def to_dict(self) -> Dict[str, Any]:
        d = {
            "rss_delta_mb": round(self.rss_delta_mb, 2),
            "rss_peak_mb": round(self.rss_peak_mb, 2),
            "gpu_delta_mb": round(self.gpu_delta_mb, 2),
            "gpu_peak_allocated_mb": round(self.gpu_peak_allocated_mb, 2),
            "gpu_peak_reserved_mb": round(self.gpu_peak_reserved_mb, 2),
            "bytes_read": self.bytes_read,
            "bytes_read_mb": round(self.bytes_read / (1024 * 1024), 2),
            "bytes_read_disk_mb": round(self.bytes_read / (1024 * 1024), 2),
            "bytes_read_logical_mb": round(self.bytes_read_logical / (1024 * 1024), 2),
            "cache_hit_ratio": round(self.cache_hit_ratio, 4),
            "bytes_written": self.bytes_written,
            "page_faults": self.page_faults,
            "major_faults": self.major_faults,
            "minor_faults": self.minor_faults,
            "working_set_estimate_mb": round(self.working_set_estimate_mb, 2),
        }
        if self.model_size_mb > 0:
            d["memory_efficiency"] = {
                "model_size_mb": round(self.model_size_mb, 2),
                "working_set_ratio": round(self.working_set_ratio, 4),
                "read_vs_size_ratio": round(self.read_vs_size_ratio, 4),
            }
        return d


def get_rss_mb() -> float:
    """
    Get current RSS (Resident Set Size) in MB.

    Reads from /proc/self/status VmRSS field.
    Returns 0.0 on failure (non-Linux or permission issues).
    """
    try:
        with open('/proc/self/status', 'r') as f:
            for line in f:
                if line.startswith('VmRSS:'):
                    parts = line.split()
                    if len(parts) >= 2:
                        return float(parts[1]) / 1024.0  # kB -> MB
    except (FileNotFoundError, PermissionError, ValueError):
        pass
    return 0.0


def get_gpu_memory_mb() -> tuple[float, float, float]:
    """
    Get GPU memory usage in MB.

    Returns: (allocated_mb, reserved_mb, peak_mb)
    Returns (0.0, 0.0, 0.0) if torch/CUDA unavailable.
    """
    try:
        import torch
        if torch.cuda.is_available():
            allocated = torch.cuda.memory_allocated() / (1024 * 1024)
            reserved = torch.cuda.memory_reserved() / (1024 * 1024)
            peak = torch.cuda.max_memory_allocated() / (1024 * 1024)
            return (allocated, reserved, peak)
    except ImportError:
        pass
    return (0.0, 0.0, 0.0)


def reset_gpu_peak_stats() -> None:
    """Reset GPU peak memory tracking."""
    try:
        import torch
        if torch.cuda.is_available():
            torch.cuda.reset_peak_memory_stats()
    except ImportError:
        pass


def get_io_bytes() -> tuple[int, int, int, int]:
    """
    Get I/O bytes from kernel counters.

    Reads from /proc/self/io.

    Returns: (read_bytes, write_bytes, rchar, wchar)
    - read_bytes: actual disk I/O (excludes cache hits)
    - write_bytes: actual disk writes
    - rchar: bytes read via read() syscalls (includes cached)
    - wchar: bytes written via write() syscalls

    WO-TRUTHY-METRICS-01: Extended to capture both disk I/O and logical I/O
    for cache-aware metrics.

    Returns (0, 0, 0, 0) on failure.
    """
    io_stats = {}
    try:
        with open("/proc/self/io") as f:
            for line in f:
                parts = line.strip().split(": ")
                if len(parts) == 2:
                    io_stats[parts[0]] = int(parts[1])
    except (FileNotFoundError, PermissionError, ValueError):
        pass
    return (
        io_stats.get("read_bytes", 0),
        io_stats.get("write_bytes", 0),
        io_stats.get("rchar", 0),
        io_stats.get("wchar", 0),
    )


def get_page_faults() -> tuple[int, int]:
    """
    Get page fault counts from getrusage().

    Returns: (major_faults, minor_faults)
    Major faults require disk I/O, minor faults are in-memory reclaims.
    """
    try:
        usage = resource.getrusage(resource.RUSAGE_SELF)
        return (usage.ru_majflt, usage.ru_minflt)
    except (AttributeError, OSError):
        return (0, 0)


@dataclass
class CacheState:
    """
    Track page cache state for cold-start accuracy.

    WO-TRUTHY-METRICS-01: Added for cache-aware benchmarking.
    """
    drop_caches_attempted: bool = False   # Whether drop_caches was called
    drop_caches_succeeded: bool = False   # Whether it succeeded (needs root)
    cold_run: bool = False                # True if this is first run after cache drop

    def to_dict(self) -> Dict[str, Any]:
        return {
            "drop_caches_attempted": self.drop_caches_attempted,
            "drop_caches_succeeded": self.drop_caches_succeeded,
            "cold_run": self.cold_run,
        }


def get_working_set_mb() -> float:
    """
    Estimate working set from /proc/self/smaps_rollup Pss.

    Pss (Proportional Set Size) accounts for shared pages more
    accurately than RSS for working set estimation.

    Returns 0.0 on failure or non-Linux.
    """
    try:
        with open('/proc/self/smaps_rollup', 'r') as f:
            for line in f:
                if line.startswith('Pss:'):
                    parts = line.split()
                    if len(parts) >= 2:
                        return float(parts[1]) / 1024.0  # kB -> MB
    except (FileNotFoundError, PermissionError, ValueError):
        pass
    return 0.0


def take_snapshot() -> MemorySnapshot:
    """
    Take a complete memory snapshot.

    Collects all available metrics (RSS, GPU, I/O, page faults, working set).
    Safe to call on non-Linux or without GPU - missing values default to 0.

    WO-TRUTHY-METRICS-01: Now captures both disk I/O (read_bytes) and logical I/O (rchar).
    """
    rss = get_rss_mb()
    gpu_alloc, gpu_reserved, gpu_peak = get_gpu_memory_mb()
    io_read, io_write, io_rchar, io_wchar = get_io_bytes()
    major_faults, minor_faults = get_page_faults()
    working_set = get_working_set_mb()

    return MemorySnapshot(
        rss_mb=rss,
        gpu_allocated_mb=gpu_alloc,
        gpu_reserved_mb=gpu_reserved,
        gpu_peak_mb=gpu_peak,
        io_read_bytes=io_read,
        io_write_bytes=io_write,
        io_rchar=io_rchar,
        io_wchar=io_wchar,
        major_faults=major_faults,
        minor_faults=minor_faults,
        working_set_mb=working_set if working_set > 0 else rss,  # Fallback to RSS
    )


def compute_delta(
    before: MemorySnapshot,
    after: MemorySnapshot,
    peak_rss_mb: Optional[float] = None,
    model_size_mb: float = 0.0,
) -> MemoryDelta:
    """
    Compute memory delta between two snapshots.

    Args:
        before: Snapshot taken before operation
        after: Snapshot taken after operation
        peak_rss_mb: Optional externally-tracked peak RSS (e.g., from loop sampling)
        model_size_mb: Optional model size for efficiency ratio calculation

    Returns:
        MemoryDelta with all changes computed

    WO-TRUTHY-METRICS-01: Now computes cache_hit_ratio and logical I/O metrics.
    """
    rss_delta = after.rss_mb - before.rss_mb
    # Peak is either externally tracked or max of before/after
    rss_peak = peak_rss_mb if peak_rss_mb is not None else max(before.rss_mb, after.rss_mb)

    gpu_delta = after.gpu_allocated_mb - before.gpu_allocated_mb
    # GPU peak comes from torch tracking
    gpu_peak_alloc = max(before.gpu_peak_mb, after.gpu_peak_mb)
    gpu_peak_reserved = max(before.gpu_reserved_mb, after.gpu_reserved_mb)

    # Disk I/O (actual reads/writes)
    bytes_read = after.io_read_bytes - before.io_read_bytes
    bytes_written = after.io_write_bytes - before.io_write_bytes

    # Logical I/O (includes cache hits) - WO-TRUTHY-METRICS-01
    bytes_read_logical = after.io_rchar - before.io_rchar
    bytes_written_logical = after.io_wchar - before.io_wchar

    # Cache hit ratio: (logical - disk) / logical
    # If logical > 0 and logical > disk, we had cache hits
    cache_hit_ratio = 0.0
    if bytes_read_logical > 0:
        cache_hits = bytes_read_logical - bytes_read
        if cache_hits >= 0:
            cache_hit_ratio = cache_hits / bytes_read_logical
        # If bytes_read > bytes_read_logical, something is off, leave ratio at 0

    major_delta = after.major_faults - before.major_faults
    minor_delta = after.minor_faults - before.minor_faults

    working_set = max(before.working_set_mb, after.working_set_mb)

    delta = MemoryDelta(
        rss_delta_mb=rss_delta,
        rss_peak_mb=rss_peak,
        gpu_delta_mb=gpu_delta,
        gpu_peak_allocated_mb=gpu_peak_alloc,
        gpu_peak_reserved_mb=gpu_peak_reserved,
        bytes_read=bytes_read,
        bytes_written=bytes_written,
        bytes_read_logical=bytes_read_logical,
        bytes_written_logical=bytes_written_logical,
        cache_hit_ratio=cache_hit_ratio,
        page_faults=major_delta + minor_delta,
        major_faults=major_delta,
        minor_faults=minor_delta,
        working_set_estimate_mb=working_set,
    )

    # Compute efficiency ratios if model size provided
    if model_size_mb > 0:
        delta.model_size_mb = model_size_mb
        delta.working_set_ratio = working_set / model_size_mb if model_size_mb > 0 else 0.0
        delta.read_vs_size_ratio = (bytes_read / (1024 * 1024)) / model_size_mb if model_size_mb > 0 else 0.0

    return delta


# Convenience class for tracking peak values during operations
class PeakTracker:
    """
    Track peak memory values during an operation.

    Usage:
        tracker = PeakTracker()
        for step in operation:
            do_step()
            tracker.sample()
        print(f"Peak RSS: {tracker.peak_rss_mb}")
    """

    def __init__(self):
        self.peak_rss_mb: float = 0.0
        self.peak_gpu_mb: float = 0.0
        self.samples: int = 0
        self._baseline_snapshot: Optional[MemorySnapshot] = None

    def start(self) -> None:
        """Initialize tracking with baseline snapshot."""
        reset_gpu_peak_stats()
        self._baseline_snapshot = take_snapshot()
        self.peak_rss_mb = self._baseline_snapshot.rss_mb
        self.peak_gpu_mb = self._baseline_snapshot.gpu_allocated_mb
        self.samples = 0

    def sample(self) -> None:
        """Sample current memory and update peaks."""
        rss = get_rss_mb()
        gpu_alloc, _, gpu_peak = get_gpu_memory_mb()

        self.peak_rss_mb = max(self.peak_rss_mb, rss)
        self.peak_gpu_mb = max(self.peak_gpu_mb, gpu_alloc, gpu_peak)
        self.samples += 1

    def finish(self) -> MemoryDelta:
        """Finish tracking and return delta from baseline."""
        after = take_snapshot()
        if self._baseline_snapshot is None:
            self._baseline_snapshot = after

        return compute_delta(
            self._baseline_snapshot,
            after,
            peak_rss_mb=self.peak_rss_mb,
        )


if __name__ == "__main__":
    # Quick self-test
    print("Memory Snapshot Utility - Self Test")
    print("=" * 50)

    snapshot = take_snapshot()
    print("Current snapshot:")
    for k, v in snapshot.to_dict().items():
        print(f"  {k}: {v}")

    print()
    print("Testing PeakTracker...")
    tracker = PeakTracker()
    tracker.start()

    # Allocate some memory
    data = [0] * (10 * 1024 * 1024)  # ~80MB of integers
    tracker.sample()

    delta = tracker.finish()
    print(f"Delta after allocation:")
    for k, v in delta.to_dict().items():
        print(f"  {k}: {v}")

    del data
    print()
    print("Self-test complete.")
