import ctypes
from ctypes import wintypes

kernel32 = ctypes.WinDLL('kernel32', use_last_error=True)

if ctypes.sizeof(ctypes.c_void_p) == 8:
    ULONG_PTR = ctypes.c_int64
else:
    ULONG_PTR = ctypes.c_ulong
DWORD = wintypes.DWORD
LARGE_INTEGER = ctypes.c_int64
SIZE_T = ULONG_PTR
ULONGLONG = ctypes.c_uint64

JobObjectAssociateCompletionPortInformation = 7
JobObjectBasicLimitInformation = 2
JobObjectBasicUIRestrictions = 4
JobObjectEndOfJobTimeInformation = 6
JobObjectExtendedLimitInformation = 9
JobObjectSecurityLimitInformation = 5
JobObjectGroupInformation = 11

JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE = 0x2000

DELETE = 0x00010000
READ_CONTROL = 0x00020000
SYNCHRONIZE = 0x00100000
WRITE_DAC = 0x00040000
WRITE_OWNER = 0x00080000
STANDARD_RIGHTS_REQUIRED = DELETE | READ_CONTROL | WRITE_DAC | WRITE_OWNER

PROCESS_CREATE_PROCESS = 0x0080
PROCESS_CREATE_THREAD = 0x0002
PROCESS_DUP_HANDLE = 0x0040
PROCESS_QUERY_INFORMATION = 0x0400
PROCESS_QUERY_LIMITED_INFORMATION = 0x1000
PROCESS_SET_INFORMATION = 0x0200
PROCESS_SET_QUOTA = 0x0100
PROCESS_SUSPEND_RESUME = 0x0800
PROCESS_TERMINATE = 0x0001
PROCESS_VM_OPERATION = 0x0008
PROCESS_VM_READ = 0x0010
PROCESS_VM_WRITE = 0x0020
PROCESS_ALL_ACCESS = STANDARD_RIGHTS_REQUIRED | SYNCHRONIZE | 0xFFF

DUPLICATE_SAME_ACCESS = 0x0002


class IO_COUNTERS(ctypes.Structure):
    _fields_ = [
        ('ReadOperationCount', ULONGLONG),
        ('WriteOperationCount', ULONGLONG),
        ('OtherOperationCount', ULONGLONG),
        ('ReadTransferCount', ULONGLONG),
        ('WriteTransferCount', ULONGLONG),
        ('OtherTransferCount', ULONGLONG),
    ]


class JOBOBJECT_BASIC_LIMIT_INFORMATION(ctypes.Structure):
    _fields_ = [
        ('PerProcessUserTimeLimit', LARGE_INTEGER),
        ('PerJobUserTimeLimit', LARGE_INTEGER),
        ('LimitFlags', DWORD),
        ('MinimumWorkingSetSize', SIZE_T),
        ('MaximumWorkingSetSize', SIZE_T),
        ('ActiveProcessLimit', DWORD),
        ('Affinity', ULONG_PTR),
        ('PriorityClass', DWORD),
        ('SchedulingClass', DWORD),
    ]


class JOBOBJECT_EXTENDED_LIMIT_INFORMATION(ctypes.Structure):
    _fields_ = [
        ('BasicLimitInformation', JOBOBJECT_BASIC_LIMIT_INFORMATION),
        ('IoInfo', IO_COUNTERS),
        ('ProcessMemoryLimit', SIZE_T),
        ('JobMemoryLimit', SIZE_T),
        ('PeakProcessMemoryUsed', SIZE_T),
        ('PeakJobMemoryUsed', SIZE_T),
    ]


class Handle(wintypes.HANDLE):
    closed = False

    def Close(self):
        if not self.closed:
            self.closed = True
            CloseHandle(self)

    def Detach(self):
        if not self.closed:
            self.closed = True
            return self.value
        raise ValueError("already closed")

    def __repr__(self):
        return "%s(%d)" % (self.__class__.__name__, self.value)

    __del__ = Close
    __str__ = __repr__


def CloseHandle(h):
    kernel32.CloseHandle(h)


def CheckError(result, msg):
    if not result:
        raise ctypes.WinError(ctypes.get_last_error(), msg)


def DuplicateHandle(
    hSourceProcess, hSourceHandle, hTargetProcess,
    desiredAccess, inheritHandle, options
):
    targetHandle = wintypes.HANDLE()
    ret = kernel32.DuplicateHandle(
        hSourceProcess, hSourceHandle, hTargetProcess,
        ctypes.byref(targetHandle), desiredAccess, inheritHandle, options)
    CheckError(ret, 'failed to duplicate handle')
    return Handle(targetHandle.value)


def GetCurrentProcess():
    hp = kernel32.GetCurrentProcess()
    return Handle(hp)


def OpenProcess(desiredAccess, inherit, pid):
    hp = kernel32.OpenProcess(desiredAccess, inherit, pid)
    CheckError(hp, 'failed to open process')
    return Handle(hp)


def CreateJobObject(jobAttributes, name):
    hp = kernel32.CreateJobObjectA(jobAttributes, name)
    CheckError(hp, 'failed to create job object')
    return Handle(hp)


def SetInformationJobObject(hJob, infoType, jobObjectInfo):
    ret = kernel32.SetInformationJobObject(
        hJob, infoType, ctypes.byref(jobObjectInfo),
        ctypes.sizeof(jobObjectInfo),
    )
    CheckError(ret, 'failed to set information job object')


def AssignProcessToJobObject(hJob, hProcess):
    ret = kernel32.AssignProcessToJobObject(hJob, hProcess)
    CheckError(ret, 'failed to assign process to job object')