# Copyright (c) 2011 Chortos-2 <chortos@inbox.lv>

from __future__ import division, with_statement

try:
	from compat import *
except ImportError:
	import __main__
	__main__.import_error(sys.exc_info()[1])

from __main__ import clock
from ctypes import *
from ctypes.wintypes import *
from subprocess import Popen

# Defaults that may be overwritten by values from _subprocess
INFINITE = -1
STD_INPUT_HANDLE = -10
WAIT_OBJECT_0 = 0

try:
	from _subprocess import *
except ImportError:
	pass

try:
	from numbers import Integral
except ImportError:
	Integral = int, long

try:
	from collections import namedtuple
except ImportError:
	from operator import itemgetter
	class ProcessTimes(tuple):
		__slots__ = ()
		__new__ = lambda cls, kernel, user: tuple.__new__(cls, (kernel, user))
		__getnewargs__ = lambda self: tuple(self)
		kernel, user = (property(itemgetter(i)) for i in (0, 1))
else:
	ProcessTimes = namedtuple('ProcessTimes', 'kernel user')

__all__ = 'call', 'kill', 'terminate'


# Automatically convert _subprocess handle objects into low-level HANDLEs
# and replicate their functionality for our own use
try:
	_subprocess_handle = type(GetCurrentProcess())
except NameError:
	_subprocess_handle = Integral
class Handle(object):
	@staticmethod
	def from_param(handle):
		if isinstance(handle, (_subprocess_handle, Integral)):
			return HANDLE(int(handle))
		elif isinstance(handle, Handle):
			return HANDLE(handle.handle)
		elif isinstance(handle, HANDLE):
			return handle
		else:
			raise TypeError('cannot convert %s to a handle' %
			                type(handle).__name__)
	
	__slots__ = 'handle'
	
	def __init__(self, handle):
		if isinstance(handle, Integral):
			self.handle = handle
		elif isinstance(handle, HANDLE):
			self.handle = handle.value
		elif isinstance(handle, Handle):
			self.handle = handle.handle
		elif isinstance(handle, _subprocess_handle):
			handle = HANDLE(int(handle))
			flags = DWORD()
			try:
				if windll.kernel32.GetHandleInformation(handle, byref(flags)):
					flags = flags.value
				else:
					flags = 0
			except AttributeError:
				# Available on NT 3.51 and up, NT line only
				flags = 0
			proc = HANDLE(int(GetCurrentProcess()))
			handle = DuplicateHandle(proc, handle, proc, 0, flags & 1, 2)
			self.handle = handle.Detach()
		else:
			raise TypeError("Handle() argument must be a handle, not '%s'" %
			                type(name).__name__)
	
	def __int__(self):
		return int(self.handle)
	
	def Detach(self):
		handle = self.handle
		self.handle = None
		return handle
	
	# This is also __del__, so only locals are accessed
	def Close(self, _CloseHandle=windll.kernel32.CloseHandle, _HANDLE=HANDLE):
		if self.handle:
			_CloseHandle(_HANDLE(self.handle))
			self.handle = None
	__del__ = Close

CHAR = c_char
INVALID_HANDLE_VALUE = HANDLE(-1).value
LPDWORD = POINTER(DWORD)
LPFILETIME = POINTER(FILETIME)
SIZE_T = ULONG_PTR = WPARAM
ULONGLONG = c_ulonglong

try:
	unicode
except NameError:
	LPCTSTR = LPCWSTR
	unisuffix = 'W'
else:
	LPCTSTR = LPCSTR
	unisuffix = 'A'


prototype = WINFUNCTYPE(BOOL, Handle,
                        LPFILETIME, LPFILETIME, LPFILETIME, LPFILETIME)
flags = ((1, 'process'),
         (2, 'creation'), (2, 'exit'), (2, 'kernel'), (2, 'user'))
try:
	GetProcessTimes = prototype(('GetProcessTimes', windll.kernel32), flags)
except AttributeError:
	# Available on NT 3.5 and up, NT line only
	GetProcessTimes = None
else:
	def errcheck(result, func, args):
		if not result: raise WinError()
		ftimes = [t.dwHighDateTime << 32 | t.dwLowDateTime for t in args[3:]]
		kernel = ftimes[0] / 10000000
		user   = ftimes[1] / 10000000
		return ProcessTimes(kernel, user)
	GetProcessTimes.errcheck = errcheck


class PROCESS_MEMORY_COUNTERS(Structure):
	_fields_ = (('cb', DWORD),
	            ('PageFaultCount', DWORD),
	            ('PeakWorkingSetSize', SIZE_T),
	            ('WorkingSetSize', SIZE_T),
	            ('QuotaPeakPagedPoolUsage', SIZE_T),
	            ('QuotaPagedPoolUsage', SIZE_T),
	            ('QuotaPeakNonPagedPoolUsage', SIZE_T),
	            ('QuotaNonPagedPoolUsage', SIZE_T),
	            ('PagefileUsage', SIZE_T),
	            ('PeakPagefileUsage', SIZE_T))

prototype = WINFUNCTYPE(BOOL, Handle, POINTER(PROCESS_MEMORY_COUNTERS), DWORD)
flags = ((1, 'process'), (2, 'counters'),
         (5, 'cb', sizeof(PROCESS_MEMORY_COUNTERS)))
try:
	GetProcessMemoryInfo = prototype(('GetProcessMemoryInfo', windll.psapi),
	                                 flags)
except AttributeError:
	# Available on NT 4.0 and up, NT line only
	GetProcessMemoryInfo = None
else:
	def errcheck(result, func, args):
		if not result: raise WinError()
		return args
	GetProcessMemoryInfo.errcheck = errcheck


class _uChar_union(Union):
	_fields_ = (('UnicodeChar', WCHAR),
	            ('AsciiChar', CHAR))

class KEY_EVENT_RECORD(Structure):
	_fields_ = (('bKeyDown', BOOL),
	            ('wRepeatCount', WORD),
	            ('wVirtualKeyCode', WORD),
	            ('wVirtualScanCode', WORD),
	            ('uChar', _uChar_union),
	            ('dwControlKeyState', DWORD))

RIGHT_ALT_PRESSED  = 0x001
LEFT_ALT_PRESSED   = 0x002
RIGHT_CTRL_PRESSED = 0x004
LEFT_CTRL_PRESSED  = 0x008
SHIFT_PRESSED      = 0x010
NUMLOCK_ON         = 0x020
SCROLLLOCK_ON      = 0x040
CAPSLOCK_ON        = 0x080
ENHANCED_KEY       = 0x100

class _Event_union(Union):
	_fields_ = ('KeyEvent', KEY_EVENT_RECORD),

class INPUT_RECORD(Structure):
	_fields_ = (('EventType', WORD),
	            ('Event', _Event_union))

KEY_EVENT                = 0x01
MOUSE_EVENT              = 0x02
WINDOW_BUFFER_SIZE_EVENT = 0x04
MENU_EVENT               = 0x08
FOCUS_EVENT              = 0x10

prototype = WINFUNCTYPE(BOOL, Handle, POINTER(INPUT_RECORD), DWORD, LPDWORD)
flags = (1, 'input'), (2, 'buffer'), (5, 'length', 1), (2, 'number_read')
ReadConsoleInput = prototype(('ReadConsoleInputA', windll.kernel32), flags)
def errcheck(result, func, args):
	if not result: raise WinError()
	return args[1] if args[3] else None
ReadConsoleInput.errcheck = errcheck


prototype = WINFUNCTYPE(BOOL, Handle)
flags = (1, 'input'),
FlushConsoleInputBuffer = prototype(('FlushConsoleInputBuffer',
                                     windll.kernel32), flags)
def errcheck(result, func, args):
	if not result: raise WinError()
FlushConsoleInputBuffer.errcheck = errcheck


prototype = WINFUNCTYPE(BOOL, Handle, DWORD)
flags = (1, 'console'), (1, 'mode')
SetConsoleMode = prototype(('SetConsoleMode', windll.kernel32), flags)
def errcheck(result, func, args):
	if not result: raise WinError()
SetConsoleMode.errcheck = errcheck

ENABLE_PROCESSED_INPUT = 0x001
ENABLE_LINE_INPUT      = 0x002
ENABLE_ECHO_INPUT      = 0x004
ENABLE_WINDOW_INPUT    = 0x008
ENABLE_MOUSE_INPUT     = 0x010
ENABLE_INSERT_MODE     = 0x020
ENABLE_QUICK_EDIT_MODE = 0x040
ENABLE_EXTENDED_FLAGS  = 0x080

ENABLE_PROCESSED_OUTPUT   = 1
ENABLE_WRAP_AT_EOL_OUTPUT = 2


prototype = WINFUNCTYPE(HANDLE, c_void_p, LPCTSTR)
flags = (5, 'attributes'), (1, 'name')
try:
	CreateJobObject = prototype(('CreateJobObject'+unisuffix, windll.kernel32),
	                            flags)
except AttributeError:
	# Available on 2000 and up, NT line only
	CreateJobObject = lambda name: None
else:
	def errcheck(result, func, args):
		if not result: raise WinError()
		return Handle(result)
	CreateJobObject.errcheck = errcheck


prototype = WINFUNCTYPE(BOOL, Handle, Handle)
flags = (1, 'job'), (1, 'handle')
try:
	AssignProcessToJobObject = prototype(('AssignProcessToJobObject',
	                                      windll.kernel32), flags)
except AttributeError:
	# Available on 2000 and up, NT line only
	AssignProcessToJobObject = lambda job, handle: None
else:
	def errcheck(result, func, args):
		if not result: raise WinError()
	AssignProcessToJobObject.errcheck = errcheck


class JOBOBJECT_BASIC_LIMIT_INFORMATION(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))

JOB_OBJECT_LIMIT_WORKINGSET                 = 0x0001
JOB_OBJECT_LIMIT_PROCESS_TIME               = 0x0002
JOB_OBJECT_LIMIT_JOB_TIME                   = 0x0004
JOB_OBJECT_LIMIT_ACTIVE_PROCESS             = 0x0008
JOB_OBJECT_LIMIT_AFFINITY                   = 0x0010
JOB_OBJECT_LIMIT_PRIORITY_CLASS             = 0x0020
JOB_OBJECT_LIMIT_PRESERVE_JOB_TIME          = 0x0040
JOB_OBJECT_LIMIT_SCHEDULING_CLASS           = 0x0080
JOB_OBJECT_LIMIT_PROCESS_MEMORY             = 0x0100
JOB_OBJECT_LIMIT_JOB_MEMORY                 = 0x0200
JOB_OBJECT_LIMIT_DIE_ON_UNHANDLED_EXCEPTION = 0x0400
JOB_OBJECT_LIMIT_BREAKAWAY_OK               = 0x0800
JOB_OBJECT_LIMIT_SILENT_BREAKAWAY_OK        = 0x1000
JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE          = 0x2000
JOB_OBJECT_LIMIT_SUBSET_AFFINITY            = 0x4000

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

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

prototype = WINFUNCTYPE(BOOL, HANDLE, c_int, c_void_p, DWORD)
flags = (1, 'job'), (1, 'infoclass'), (1, 'info'), (1, 'infosize')
try:
	_setjobinfo = prototype(('SetInformationJobObject',windll.kernel32), flags)
except AttributeError:
	# Available on 2000 and up, NT line only
	SetInformationJobObject = lambda job, infoclass, info: None
else:
	def errcheck(result, func, args):
		if not result: raise WinError()
	_setjobinfo.errcheck = errcheck
	def SetInformationJobObject(job, infoclass, info):
		return _setjobinfo(job, infoclass, info, sizeof(info))

(
	JobObjectBasicAccountingInformation,
	JobObjectBasicLimitInformation,
	JobObjectBasicProcessIdList,
	JobObjectBasicUIRestrictions,
	JobObjectSecurityLimitInformation,   
	JobObjectEndOfJobTimeInformation,
	JobObjectAssociateCompletionPortInformation,
	JobObjectBasicAndIoAccountingInformation,
	JobObjectExtendedLimitInformation,
	JobObjectJobSetInformation,
	MaxJobObjectInfoClass
) = range(1, 12)


prototype = WINFUNCTYPE(DWORD, DWORD, POINTER(HANDLE), BOOL, DWORD)
flags = (1, 'count'), (1, 'handles'), (1, 'wait_all'), (1, 'milliseconds')
_wait_multiple = prototype(('WaitForMultipleObjects', windll.kernel32), flags)
def errcheck(result, func, args):
	if result == WAIT_FAILED: raise WinError()
	return args
_wait_multiple.errcheck = errcheck
def WaitForMultipleObjects(handles, wait_all, timeout):
	n = len(handles)
	handles = (Handle.from_param(handle) for handle in handles)
	timeout = ceil(timeout * 1000)
	return _wait_multiple(n, (HANDLE * n)(*handles), wait_all, timeout)

# WAIT_OBJECT_0 defined at the top of the file
WAIT_ABANDONED_0 = 0x00000080
WAIT_TIMEOUT     = 0x00000102
WAIT_FAILED      = 0xFFFFFFFF


try:
	_wait_single = WaitForSingleObject
except NameError:
	prototype = WINFUNCTYPE(DWORD, Handle, DWORD)
	flags = (1, 'handle'), (1, 'milliseconds')
	_wait_single = prototype(('WaitForSingleObject', windll.kernel32), flags)
	def errcheck(result, func, args):
		if result == WAIT_FAILED: raise WinError()
		return args
	_wait_single.errcheck = errcheck
def WaitForSingleObject(handle, timeout):
	return _wait_single(handle, ceil(timeout * 1000))


try:
	GetStdHandle
except NameError:
	prototype = WINFUNCTYPE(HANDLE, DWORD)
	flags = (1, 'which'),
	GetStdHandle = prototype(('GetStdHandle', windll.kernel32), flags)
	def errcheck(result, func, args):
		if result == INVALID_HANDLE_VALUE: raise WinError()
		return args if result else None
	GetStdHandle.errcheck = errcheck


try:
	TerminateProcess
except NameError:
	prototype = WINFUNCTYPE(BOOL, Handle, UINT)
	flags = (1, 'process'), (1, 'exitcode')
	TerminateProcess = prototype(('TerminateProcess', windll.kernel32), flags)
	def errcheck(result, func, args):
		if not result: raise WinError()
	TerminateProcess.errcheck = errcheck


# Do not show error messages due to errors in the program being tested
try:
	errmode = ctypes.windll.kernel32.GetErrorMode()
except AttributeError:
	# GetErrorMode is available on Vista/2008 and up
	errmode = ctypes.windll.kernel32.SetErrorMode(0)
ctypes.windll.kernel32.SetErrorMode(errmode | 0x8003)

stdin = GetStdHandle(STD_INPUT_HANDLE)
try:
	SetConsoleMode(stdin, ENABLE_PROCESSED_INPUT)
except WindowsError:
	console_input = False
else:
	console_input = True
	FlushConsoleInputBuffer(stdin)

def kill(process):
	try:
		process.terminate()
	except AttributeError:
		TerminateProcess(process._handle)
terminate = kill

def call(*args, **kwargs):
	case = kwargs.pop('case')
	job = CreateJobObject(None)
	flags = 0
	if case.maxcputime:
		flags |= JOB_OBJECT_LIMIT_PROCESS_TIME
	if case.maxmemory:
		flags |= JOB_OBJECT_LIMIT_PROCESS_MEMORY
	limits = JOBOBJECT_EXTENDED_LIMIT_INFORMATION(
		JOBOBJECT_BASIC_LIMIT_INFORMATION(
			PerProcessUserTimeLimit=ceil((case.maxcputime or 0)*10000000),
			LimitFlags=flags,
		),
		ProcessMemoryLimit=ceil((case.maxmemory or 0)*1048576),
	)
	SetInformationJobObject(job, JobObjectExtendedLimitInformation, limits)
	try:
		case.process = Popen(*args, **kwargs)
	except OSError:
		raise CannotStartTestee(sys.exc_info()[1])
	case.time_started = clock()
	AssignProcessToJobObject(job, case.process._handle)
	if not console_input:
		if case.maxwalltime:
			if (WaitForSingleObject(case.process._handle, case.maxwalltime) !=
			    WAIT_OBJECT_0):
				raise TimeLimitExceeded
		else:
			case.process.wait()
	else:
		handles = stdin, case.process._handle
		if case.maxwalltime:
			time_end = clock() + case.maxwalltime
			while case.process.poll() is None:
				remaining = time_end - clock()
				if remaining > 0:
					if (WaitForMultipleObjects(handles, False, remaining) ==
					    WAIT_OBJECT_0):
						ir = ReadConsoleInput(stdin)
						if (ir and
						    ir.EventType == 1 and
						    ir.Event.KeyEvent.bKeyDown and
						    ir.Event.KeyEvent.wVirtualKeyCode == 27):
							raise CanceledByUser
				else:
					raise TimeLimitExceeded
		else:
			while case.process.poll() is None:
				if (WaitForMultipleObjects(handles, False, INFINITE) ==
				    WAIT_OBJECT_0):
					ir = ReadConsoleInput(stdin)
					if (ir and
					    ir.EventType == 1 and
					    ir.Event.KeyEvent.bKeyDown and
					    ir.Event.KeyEvent.wVirtualKeyCode == 27):
						raise CanceledByUser
	case.time_stopped = clock()
	if case.maxcputime and GetProcessTimes:
		try:
			times = GetProcessTimes(case.process._handle)
		except WindowsError:
			pass
		else:
			if times.kernel + times.user > case.maxcputime:
				raise TimeLimitExceeded
	if case.maxmemory and GetProcessMemoryInfo:
		try:
			counters = GetProcessMemoryInfo(case.process._handle)
		except WindowsError:
			pass
		else:
			if counters.PeakPagefileUsage > case.maxmemory * 1048576:
				raise MemoryLimitExceeded