"""The kernel's memory cap: a sampled physical footprint, and an allocation limit.

A cell that pushes the process past the cap is interrupted within an
interval. One allocation can still jump past a sample, so the kernel also
sets a soft address-space limit above the cap (`enforce`): an allocation past
it raises MemoryError inside the cell. It is a guard against a runaway cell
taking the namespace down, not a sandbox.
"""

from __future__ import annotations

import builtins
import contextlib
import ctypes
import io
import os
import struct
import subprocess
import sys
import threading
from collections.abc import Generator, Iterator
from types import TracebackType
from typing import IO, Any

LIMIT = int(os.environ.get("ALBEDO_KERNEL_MEMORY_BYTES", str(256 * 1024 * 1024)))
INTERVAL = 0.1  # seconds between samples while a cell runs
# A cell that starts with the kernel already past the cap (a large variable
# left behind) may still grow this much, so `del big` and gc.collect() run.
SLACK = 8 * 1024 * 1024
PAGE = os.sysconf("SC_PAGE_SIZE") if hasattr(os, "sysconf") else 4096
# What the soft limit allows above the cap. Linux counts private writable
# mappings only (thread stacks included); Darwin counts the whole address
# space, file mappings too, and its allocator keeps freed regions mapped, so
# it gets room for those and the limit is a fuse, not a soft cap.
MARGIN = LIMIT // 2 if sys.platform == "linux" else LIMIT * 2
# Extra room while a cell that hit the limit is reported.
HEADROOM = 64 * 1024 * 1024
# The most a read of a large file may take in one call.
READ_CHUNK = max(1024 * 1024, LIMIT // 8)


def footprint() -> int | None:
    """Bytes this process holds in physical memory: the resident set from
    /proc, or Darwin's physical footprint (what vmmap and Activity Monitor
    report). None where neither can tell."""
    try:
        with open("/proc/self/statm", "rb") as statm:
            return int(statm.read().split()[1]) * PAGE
    except (OSError, IndexError, ValueError):
        return darwin_footprint(os.getpid()) if sys.platform == "darwin" else None


def darwin_footprint(pid: int) -> int | None:
    """ri_phys_footprint of a Darwin process's rusage_info_v0 (libproc's
    proc_pid_rusage), the 10th field after the 16-byte uuid."""
    try:
        libproc = ctypes.CDLL("/usr/lib/libproc.dylib")
    except OSError:
        return None
    buffer = ctypes.create_string_buffer(96)
    if libproc.proc_pid_rusage(pid, 0, buffer) != 0:
        return None
    return struct.unpack_from("=Q", buffer.raw, 72)[0]


def bloated() -> bool:
    """Whether the kernel holds more than its cap right now."""
    used = footprint()
    return used is not None and used > LIMIT


def threshold(start: int) -> int:
    """The footprint past which a cell that started at `start` is interrupted."""
    return max(LIMIT, start + SLACK)


def mebibytes(size: int) -> str:
    return f"{size / (1024 * 1024):.0f} MiB"


def address_space() -> int | None:
    """The size the soft limit measures: Linux's VmData (private writable
    mappings), Darwin's virtual size. None where neither can tell."""
    if sys.platform == "linux":
        try:
            with open("/proc/self/status", "rb") as status:
                for line in status:
                    if line.startswith(b"VmData:"):
                        return int(line.split()[1]) * 1024
        except (OSError, IndexError, ValueError):
            return None
    if sys.platform == "darwin":
        try:
            libc = ctypes.CDLL(None)
            libc.mach_task_self.restype = ctypes.c_uint32
            info = (ctypes.c_uint32 * 256)()
            count = ctypes.c_uint32(256)
            # TASK_VM_INFO: virtual_size is its first field.
            if libc.task_info(libc.mach_task_self(), 22, info, ctypes.byref(count)):
                return None
            return ctypes.cast(info, ctypes.POINTER(ctypes.c_uint64))[0]
        except (OSError, AttributeError):
            return None
    return None


def limit_resource() -> int | None:
    import resource

    if sys.platform == "linux":
        return resource.RLIMIT_DATA
    if sys.platform == "darwin":
        return resource.RLIMIT_AS
    return None


_installed: int | None = None  # the soft limit enforce() set
_lifts = 0
_lock = threading.RLock()


def _set_soft(soft: int) -> None:
    import resource

    res = limit_resource()
    if res is not None:
        resource.setrlimit(res, (soft, resource.getrlimit(res)[1]))


def enforce() -> int | None:
    """Sets the soft limit at the current address space plus the cap and
    MARGIN, so an allocation past it raises MemoryError in the cell. Returns
    the limit, or None where the platform or the environment refuses one."""
    global _installed
    base = address_space()
    if base is None or LIMIT <= 0:
        return None
    import resource

    res = limit_resource()
    assert res is not None
    soft, hard = resource.getrlimit(res)
    target = base + LIMIT + MARGIN
    if hard != resource.RLIM_INFINITY and target > hard:
        return None
    try:
        resource.setrlimit(res, (target, hard))
    except (ValueError, OSError):
        return None
    with _lock:
        _installed = target
    return target


@contextlib.contextmanager
def lifted() -> Generator[None]:
    """The limit off, for a call that starts a process: children inherit
    the soft limit, and a build or a model they run is not the kernel's."""
    global _lifts
    import resource

    res = limit_resource()
    if _installed is None or res is None:
        yield
        return
    with _lock:
        if _lifts == 0:
            _set_soft(resource.getrlimit(res)[1])
        _lifts += 1
    try:
        yield
    finally:
        with _lock:
            _lifts -= 1
            if _lifts == 0 and _installed is not None:
                _set_soft(_installed)


@contextlib.contextmanager
def headroom() -> Generator[None]:
    """Room above the limit while a cell that hit it is reported, so the
    kernel can still build its reply."""
    global _installed
    if _installed is None:
        yield
        return
    with _lock:
        previous = _installed
        _installed = previous + HEADROOM
        if _lifts == 0:
            _set_soft(_installed)
    try:
        yield
    finally:
        with _lock:
            _installed = previous
            if _lifts == 0:
                _set_soft(previous)


def refusal(path: object, size: int) -> MemoryError:
    return MemoryError(
        f"{path} is {mebibytes(size)}; reading it whole would pass the kernel's"
        f" {mebibytes(LIMIT)} memory cap. Page it with files.read(path, start_line=,"
        " end_line=), iterate the file line by line, or read(n) in chunks."
    )


class LargeFile:
    """A file past the cap: read() without a size, and readlines(), refuse."""

    def __init__(self, file: IO[Any], path: object, size: int) -> None:
        self._file, self._path, self._size = file, path, size

    def __getattr__(self, name: str) -> object:
        return getattr(self._file, name)

    def __iter__(self) -> Iterator[object]:
        return iter(self._file)

    def __enter__(self) -> "LargeFile":
        _ = self._file.__enter__()
        return self

    def __exit__(
        self,
        kind: type[BaseException] | None,
        error: BaseException | None,
        trace: TracebackType | None,
    ) -> None:
        self._file.__exit__(kind, error, trace)

    def read(self, size: int | None = -1) -> object:
        if size is None or size < 0 or size > READ_CHUNK:
            raise refusal(self._path, self._size)
        return self._file.read(size)

    def readlines(self, hint: int = -1) -> object:
        if hint is None or hint <= 0 or hint > READ_CHUNK:
            raise refusal(self._path, self._size)
        return self._file.readlines(hint)


def guard_open(real):  # type: ignore[no-untyped-def]
    """open() that returns a LargeFile for a file past the cap opened for
    reading; every other file is returned as it is."""

    def guarded(file, mode="r", *args, **kwargs):  # type: ignore[no-untyped-def]
        opened = real(file, mode, *args, **kwargs)
        if isinstance(file, int) or "r" not in mode or "+" in mode:
            return opened
        try:
            size = os.fstat(opened.fileno()).st_size
            regular = os.path.isfile(os.fspath(file))
        except (OSError, TypeError, ValueError):
            return opened
        if regular and size > LIMIT:
            return LargeFile(opened, file, size)
        return opened

    guarded.__name__ = getattr(real, "__name__", "open")
    return guarded


def lift_around(real):  # type: ignore[no-untyped-def]
    def wrapped(*args, **kwargs):  # type: ignore[no-untyped-def]
        with lifted():
            return real(*args, **kwargs)

    wrapped.__name__ = getattr(real, "__name__", "spawn")
    return wrapped


def install() -> int | None:
    """Starts enforcing: the soft limit, and the calls that must not hand it
    to a child or load a file past the cap whole. Returns the limit."""
    limit = enforce()
    if limit is not None:
        setattr(
            subprocess.Popen,
            "_execute_child",
            lift_around(getattr(subprocess.Popen, "_execute_child")),
        )
        for name in ("system", "posix_spawn", "posix_spawnp"):
            if hasattr(os, name):
                setattr(os, name, lift_around(getattr(os, name)))
    builtins.open = guard_open(builtins.open)
    io.open = builtins.open  # type: ignore[assignment]
    return limit