import psutil
from dataclasses import dataclass
import pandas as pd


@dataclass
class MemoryGuard:
    """Prevent OOM by checking memory before processing"""

    max_memory_gb: float = 8.0
    safety_factor: float = 0.7

    def check_feasibility(self, df: pd.DataFrame, chunk_size: int) -> tuple[bool, str]:
        """Check if processing is feasible with current memory"""
        # Estimate memory usage
        avg_str_len = 50  # Conservative estimate
        n_string_cols = len(df.select_dtypes(include=["object", "string"]).columns)

        # Bytes per chunk
        chunk_memory = chunk_size * n_string_cols * avg_str_len

        # Add overhead for transformations (2x for copies)
        estimated_usage = chunk_memory * 2

        # Available memory
        available = psutil.virtual_memory().available
        max_allowed = self.max_memory_gb * 1024 * 1024 * 1024 * self.safety_factor

        if estimated_usage > min(available, max_allowed):
            return False, (
                f"Estimated memory usage ({estimated_usage / 1e9:.1f}GB) exceeds "
                f"available ({available / 1e9:.1f}GB) or limit ({max_allowed / 1e9:.1f}GB). "
                f"Reduce chunk_size or increase memory."
            )

        return True, "OK"
