import ctypes
import gc
import logging
import time


class MemoryCleanupManager:
    """Schedule and execute periodic memory cleanup for an agent.

    This manager provides independent frame-based and time-based triggers for
    Python's garbage collector and glibc's malloc_trim (Linux only). It is
    designed to be lightweight and safe to call from a frame-complete hook.
    """

    def __init__(
        self,
        logger: logging.Logger,
        *,
        gc_every_n_frames: int = 0,
        gc_every_n_seconds: float = 0.0,
        malloc_trim_every_n_frames: int = 0,
        malloc_trim_every_n_seconds: float = 0.0,
    ):
        """Initialize the cleanup scheduler.

        Args:
            logger: Logger for warnings and timing information.
            gc_every_n_frames: Run gc.collect every N frames (0 disables).
            gc_every_n_seconds: Run gc.collect every N seconds (0 disables).
            malloc_trim_every_n_frames: Run malloc_trim every N frames (0 disables).
            malloc_trim_every_n_seconds: Run malloc_trim every N seconds (0 disables).
        """
        self.logger = logger
        self.gc_every_n_frames = max(0, int(gc_every_n_frames))
        self.gc_every_n_seconds = max(0.0, float(gc_every_n_seconds))
        self.malloc_trim_every_n_frames = max(0, int(malloc_trim_every_n_frames))
        self.malloc_trim_every_n_seconds = max(0.0, float(malloc_trim_every_n_seconds))

        now = time.monotonic()
        self._gc_frames_since = 0
        self._gc_last_ts = now
        self._trim_frames_since = 0
        self._trim_last_ts = now
        self._malloc_trim_fn = None
        self._malloc_trim_unavailable = False

    def _enabled(self) -> bool:
        """Return True if any cleanup trigger is enabled."""
        return any(
            value > 0
            for value in (
                self.gc_every_n_frames,
                self.gc_every_n_seconds,
                self.malloc_trim_every_n_frames,
                self.malloc_trim_every_n_seconds,
            )
        )

    def on_frame_complete(self) -> None:
        """Update counters and run any cleanup actions that are due."""
        if not self._enabled():
            return

        now = time.monotonic()
        self._gc_frames_since += 1
        self._trim_frames_since += 1

        run_gc = self._should_run(
            self._gc_frames_since,
            self.gc_every_n_frames,
            self._gc_last_ts,
            self.gc_every_n_seconds,
            now,
        )
        run_trim = self._should_run(
            self._trim_frames_since,
            self.malloc_trim_every_n_frames,
            self._trim_last_ts,
            self.malloc_trim_every_n_seconds,
            now,
        )

        if run_gc:
            self._run_gc(now)
        if run_trim:
            self._run_malloc_trim(now)

    def _should_run(
        self,
        frames_since: int,
        every_n_frames: int,
        last_ts: float,
        every_n_seconds: float,
        now: float,
    ) -> bool:
        """Return True if frame or time thresholds indicate a run is due."""
        if every_n_frames > 0 and frames_since >= every_n_frames:
            return True
        if every_n_seconds > 0.0 and (now - last_ts) >= every_n_seconds:
            return True
        return False

    def _run_gc(self, now: float) -> None:
        """Run gc.collect() and log if it takes longer than 100ms."""
        start = time.monotonic()
        try:
            gc.collect()
        except Exception as exc:
            self.logger.warning("Memory cleanup: gc.collect failed: %s", exc)
        elapsed = time.monotonic() - start
        if elapsed > 0.1:
            self.logger.info("Memory cleanup: gc.collect took %.3fs", elapsed)
        self._gc_frames_since = 0
        self._gc_last_ts = now

    def _ensure_malloc_trim(self) -> bool:
        """Load libc and cache malloc_trim if available; return True if usable."""
        if self._malloc_trim_unavailable:
            return False
        if self._malloc_trim_fn is not None:
            return True
        try:
            libc = ctypes.CDLL("libc.so.6")
            malloc_trim = libc.malloc_trim
            malloc_trim.argtypes = [ctypes.c_size_t]
            malloc_trim.restype = ctypes.c_int
            self._malloc_trim_fn = malloc_trim
            return True
        except Exception as exc:
            self._malloc_trim_unavailable = True
            self.logger.info("Memory cleanup: malloc_trim unavailable: %s", exc)
            self.malloc_trim_every_n_frames = 0
            self.malloc_trim_every_n_seconds = 0.0
            return False

    def _run_malloc_trim(self, now: float) -> None:
        """Run malloc_trim(0) and log if it takes longer than 100ms."""
        if self._ensure_malloc_trim():
            start = time.monotonic()
            try:
                self._malloc_trim_fn(0)
            except Exception as exc:
                self.logger.warning("Memory cleanup: malloc_trim failed: %s", exc)
            elapsed = time.monotonic() - start
            if elapsed > 0.1:
                self.logger.info("Memory cleanup: malloc_trim took %.3fs", elapsed)

        self._trim_frames_since = 0
        self._trim_last_ts = now
