import os
import gc
import psutil
from typing import Optional, Any, Callable

def cleanup_memory(
    gen: int = 0,
    tag: str = "",
    ctx: Any = None,
    logger: Any = None,
    mem_hiwater_mb: Optional[int] = None,
    restart_executor_fn: Optional[Callable[[int], None]] = None,
    restart_every_n: int = 50,
):
    """
    Aggressively release memory and (optionally) rotate the executor.
    
    This utility is designed to be called by the GA engine to manage 
    resource consumption during long runs.
    """
    # 1. Determine current memory usage
    try:
        process = psutil.Process(os.getpid())
        mem_mb = process.memory_info().rss / (1024 ** 2)
    except Exception:
        process, mem_mb = None, None

    # 2. Check triggers: cadence OR high-water mark
    meets_cadence = (gen % 10 == 0)
    meets_hiwater = (mem_mb is not None and mem_hiwater_mb is not None and mem_mb >= mem_hiwater_mb)

    if not (meets_cadence or meets_hiwater):
        return True

    if logger:
        logger.info(f"[Memory cleanup] Begin {tag} (gen={gen}, mem={mem_mb:.1f}MB, hiwater={mem_hiwater_mb}MB)...")

    # 3. Clear temporary context data
    if ctx:
        if hasattr(ctx, "clear_temporary"):
            ctx.clear_temporary()
        else:
            for attr in ("_features", "_objectives", "_selection", "top_objectives"):
                if hasattr(ctx, attr):
                    setattr(ctx, attr, None)

    # 4. Force Garbage Collection
    collected = gc.collect()

    # 5. Measure after
    try:
        after = process.memory_info().rss / (1024 ** 2) if process else None
    except Exception:
        after = None

    if logger:
        if mem_mb is not None and after is not None:
            logger.info(
                "[Memory cleanup] Done %s: %d objects collected, freed %+0.1f MB (now %.1f MB).",
                tag, collected, mem_mb - after, after
            )
        else:
            logger.info("[Memory cleanup] Done %s: %d objects collected.", tag, collected)

    # 6. Optional: restart executor (safe & conditional)
    if restart_executor_fn:
        try:
            meets_restart_cadence = (restart_every_n > 0 and gen % restart_every_n == 0)
            # Re-check hiwater with post-GC memory
            meets_hiwater_after = (after is not None and mem_hiwater_mb is not None and after >= mem_hiwater_mb)
            
            if meets_restart_cadence or meets_hiwater_after:
                # Note: Safety check (if job is running) should be handled by the engine 
                # or passed via a more complex status object. 
                # For now, we assume if restart_executor_fn is called, it might rotate.
                restart_executor_fn(gen)
        except Exception as e:
            if logger:
                logger.debug("Executor restart step failed: %s", e)

    return True
