"""Resource guards: keep the search inside the machine's RAM and disk.

Nothing here touches the GPU; the whole project is CPU and numpy only.
"""
from __future__ import annotations

import ctypes
import os
import shutil
import sys
from pathlib import Path

MB = 1024 * 1024
GB = 1024 * MB

_JOB_HANDLE = None  # kept alive so the Windows job object is not released


def free_disk_gb(path: str | Path) -> float:
    return shutil.disk_usage(str(path)).free / GB


def dir_size_mb(path: str | Path) -> float:
    total = 0
    for root, _dirs, files in os.walk(str(path)):
        for name in files:
            try:
                total += os.path.getsize(os.path.join(root, name))
            except OSError:
                continue
    return total / MB


def free_ram_mb() -> float | None:
    """Physical RAM currently available, or None if it cannot be read."""
    if sys.platform == "win32":
        return _free_ram_mb_windows()
    try:
        with open("/proc/meminfo", encoding="ascii") as handle:
            for line in handle:
                if line.startswith("MemAvailable:"):
                    return int(line.split()[1]) / 1024
    except OSError:
        return None
    return None


def _free_ram_mb_windows() -> float | None:
    class MemoryStatus(ctypes.Structure):
        _fields_ = [
            ("dwLength", ctypes.c_uint32),
            ("dwMemoryLoad", ctypes.c_uint32),
            ("ullTotalPhys", ctypes.c_uint64),
            ("ullAvailPhys", ctypes.c_uint64),
            ("ullTotalPageFile", ctypes.c_uint64),
            ("ullAvailPageFile", ctypes.c_uint64),
            ("ullTotalVirtual", ctypes.c_uint64),
            ("ullAvailVirtual", ctypes.c_uint64),
            ("ullAvailExtendedVirtual", ctypes.c_uint64),
        ]

    status = MemoryStatus()
    status.dwLength = ctypes.sizeof(MemoryStatus)
    if not ctypes.windll.kernel32.GlobalMemoryStatusEx(ctypes.byref(status)):
        return None
    return status.ullAvailPhys / MB


def confine_own_process(limit_mb: int) -> bool:
    """Hard-cap this process's memory; on Windows also forbid it from starting processes.

    Allocations past the cap fail with MemoryError. Returns True when the cap is in force.
    """
    if sys.platform == "win32":
        return _cap_memory_windows(limit_mb)
    try:
        import resource

        limit = limit_mb * MB
        resource.setrlimit(resource.RLIMIT_AS, (limit, limit))
        return True
    except (ImportError, ValueError, OSError):
        return False


def _cap_memory_windows(limit_mb: int) -> bool:
    global _JOB_HANDLE
    from ctypes import wintypes

    class BasicLimits(ctypes.Structure):
        _fields_ = [
            ("PerProcessUserTimeLimit", ctypes.c_int64),
            ("PerJobUserTimeLimit", ctypes.c_int64),
            ("LimitFlags", wintypes.DWORD),
            ("MinimumWorkingSetSize", ctypes.c_size_t),
            ("MaximumWorkingSetSize", ctypes.c_size_t),
            ("ActiveProcessLimit", wintypes.DWORD),
            ("Affinity", ctypes.c_size_t),
            ("PriorityClass", wintypes.DWORD),
            ("SchedulingClass", wintypes.DWORD),
        ]

    class IoCounters(ctypes.Structure):
        _fields_ = [(name, ctypes.c_uint64) for name in (
            "ReadOperationCount", "WriteOperationCount", "OtherOperationCount",
            "ReadTransferCount", "WriteTransferCount", "OtherTransferCount")]

    class ExtendedLimits(ctypes.Structure):
        _fields_ = [
            ("BasicLimitInformation", BasicLimits),
            ("IoInfo", IoCounters),
            ("ProcessMemoryLimit", ctypes.c_size_t),
            ("JobMemoryLimit", ctypes.c_size_t),
            ("PeakProcessMemoryUsed", ctypes.c_size_t),
            ("PeakJobMemoryUsed", ctypes.c_size_t),
        ]

    job_object_extended_limit_information = 9
    job_object_limit_process_memory = 0x00000100
    job_object_limit_active_process = 0x00000008
    job_object_limit_kill_on_job_close = 0x00002000

    k32 = ctypes.WinDLL("kernel32", use_last_error=True)
    k32.CreateJobObjectW.restype = wintypes.HANDLE
    k32.CreateJobObjectW.argtypes = [ctypes.c_void_p, wintypes.LPCWSTR]
    k32.GetCurrentProcess.restype = wintypes.HANDLE
    k32.SetInformationJobObject.restype = wintypes.BOOL
    k32.SetInformationJobObject.argtypes = [
        wintypes.HANDLE, ctypes.c_int, ctypes.c_void_p, wintypes.DWORD]
    k32.AssignProcessToJobObject.restype = wintypes.BOOL
    k32.AssignProcessToJobObject.argtypes = [wintypes.HANDLE, wintypes.HANDLE]

    job = k32.CreateJobObjectW(None, None)
    if not job:
        return False
    info = ExtendedLimits()
    info.BasicLimitInformation.LimitFlags = (
        job_object_limit_process_memory | job_object_limit_active_process
        | job_object_limit_kill_on_job_close)
    info.BasicLimitInformation.ActiveProcessLimit = 1
    info.ProcessMemoryLimit = limit_mb * MB
    if not k32.SetInformationJobObject(
            job, job_object_extended_limit_information,
            ctypes.byref(info), ctypes.sizeof(info)):
        return False
    if not k32.AssignProcessToJobObject(job, k32.GetCurrentProcess()):
        return False
    _JOB_HANDLE = job
    return True


def check_budget(run_dir: str | Path, parallel: int, mem_mb: int,
                 min_free_disk_gb: float, max_run_mb: float) -> str | None:
    """Return a reason to stop, or None when the run is inside its budget."""
    free_gb = free_disk_gb(run_dir)
    if free_gb < min_free_disk_gb:
        return f"only {free_gb:.1f} GB disk free (floor is {min_free_disk_gb} GB)"
    used_mb = dir_size_mb(run_dir)
    if used_mb > max_run_mb:
        return f"run folder is {used_mb:.0f} MB (cap is {max_run_mb:.0f} MB)"
    free_mb = free_ram_mb()
    needed = parallel * mem_mb
    if free_mb is not None and needed > 0.5 * free_mb:
        return (f"{parallel} workers x {mem_mb} MB = {needed} MB, more than half of the "
                f"{free_mb:.0f} MB RAM currently free; lower --parallel or --mem-mb")
    return None