"""Windows resource containment for the inspector and all descendant workers.

This is a Job Object, not a restricted-token or filesystem/network sandbox.
The inspector owns the handle: terminating it also terminates its model workers.
"""
import ctypes
import sys
from ctypes import wintypes

_job_handle = None
MEMORY_LIMIT = 4 * 1024 * 1024 * 1024
PROCESS_LIMIT = 4


def contain_inspector():
    global _job_handle
    if sys.platform != "win32":
        return {"active": False, "kind": "unsupported-platform"}
    if _job_handle:
        return {"active": True, "kind": "windows-job-v1"}

    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)]

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

    handle = kernel.CreateJobObjectW(None, None)
    if not handle:
        raise RuntimeError("Inspector resource containment could not be created")
    limits = ExtendedLimits()
    # KILL_ON_JOB_CLOSE | JOB_MEMORY | ACTIVE_PROCESS. No BREAKAWAY flags.
    limits.BasicLimitInformation.LimitFlags = 0x2000 | 0x200 | 0x8
    limits.BasicLimitInformation.ActiveProcessLimit = PROCESS_LIMIT
    limits.JobMemoryLimit = MEMORY_LIMIT
    if not kernel.SetInformationJobObject(handle, 9, ctypes.byref(limits), ctypes.sizeof(limits)):
        kernel.CloseHandle(handle)
        raise RuntimeError("Inspector resource containment limits could not be applied")
    if not kernel.AssignProcessToJobObject(handle, kernel.GetCurrentProcess()):
        kernel.CloseHandle(handle)
        raise RuntimeError("Inspector resource containment could not be attached")
    # Retain a non-inheritable owner handle until process exit. Do not close early:
    # KILL_ON_JOB_CLOSE also terminates the current inspector process.
    _job_handle = handle
    return {"active": True, "kind": "windows-job-v1", "memoryBytes": MEMORY_LIMIT,
            "maxProcesses": PROCESS_LIMIT, "killsDescendants": True}