"""macOS process memory accounting, including Metal allocations.

Physical footprint, resident size and MLX allocation counters overlap. They
must be reported separately. Sampling can miss short peaks; the lifetime
high-water counter includes all earlier requests in a persistent process.
"""

from __future__ import annotations

import ctypes
import os
import threading


class _RUsageV4(ctypes.Structure):
    # Exact macOS SDK sys/resource.h rusage_info_v4 layout.
    _fields_ = [("uuid", ctypes.c_uint8 * 16)] + [
        (name, ctypes.c_uint64)
        for name in (
            "user_time",
            "system_time",
            "pkg_idle_wkups",
            "interrupt_wkups",
            "pageins",
            "wired_size",
            "resident_size",
            "phys_footprint",
            "proc_start_abstime",
            "proc_exit_abstime",
            "child_user_time",
            "child_system_time",
            "child_pkg_idle_wkups",
            "child_interrupt_wkups",
            "child_pageins",
            "child_elapsed_abstime",
            "diskio_bytesread",
            "diskio_byteswritten",
            "cpu_time_qos_default",
            "cpu_time_qos_maintenance",
            "cpu_time_qos_background",
            "cpu_time_qos_utility",
            "cpu_time_qos_legacy",
            "cpu_time_qos_user_initiated",
            "cpu_time_qos_user_interactive",
            "billed_system_time",
            "serviced_system_time",
            "logical_writes",
            "lifetime_max_phys_footprint",
            "instructions",
            "cycles",
            "billed_energy",
            "serviced_energy",
            "interval_max_phys_footprint",
            "runnable_time",
        )
    ]


_proc = ctypes.CDLL("/usr/lib/libproc.dylib", use_errno=True)
_proc.proc_pid_rusage.argtypes = [ctypes.c_int, ctypes.c_int, ctypes.c_void_p]
_proc.proc_pid_rusage.restype = ctypes.c_int


def footprint():
    info = _RUsageV4()
    if _proc.proc_pid_rusage(os.getpid(), 4, ctypes.byref(info)):
        raise OSError(ctypes.get_errno(), "proc_pid_rusage")
    return {
        "physical_bytes": info.phys_footprint,
        "lifetime_peak_physical_bytes": info.lifetime_max_phys_footprint,
        "resident_bytes": info.resident_size,
    }


class FootprintMonitor:
    """Sample current footprint; lifetime high water includes earlier runs."""

    def __init__(self, interval=0.1):
        self.interval = interval
        self.samples = []
        self.stop_event = threading.Event()

    def _run(self):
        while not self.stop_event.wait(self.interval):
            self.samples.append(footprint())

    def __enter__(self):
        self.samples.append(footprint())
        self.worker = threading.Thread(target=self._run, daemon=True)
        self.worker.start()
        return self

    def __exit__(self, *_):
        self.stop_event.set()
        self.worker.join()
        self.samples.append(footprint())

    def result(self):
        return {
            "sampled_peak_physical_bytes": max(
                s["physical_bytes"] for s in self.samples
            ),
            "lifetime_peak_physical_bytes": self.samples[-1][
                "lifetime_peak_physical_bytes"
            ],
            "sample_interval_s": self.interval,
            "sample_count": len(self.samples),
            "scope": "Parent process only; includes Metal; excludes compiler XPC services. Counters overlap; do not sum.",
        }