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

from __future__ import division, with_statement

from compat import *
import testcases  # mutual import

from subprocess import Popen
import os, sys, time

try:
	from testcases import clock
except ImportError:
	if sys.platform.startswith('java'):
		from time import clock
	else:
		from time import time as clock

try:
	from signal import SIGTERM, SIGKILL
except ImportError:
	SIGTERM = 15
	SIGKILL = 9

__all__ = 'call', 'kill', 'terminate', 'pause', 'clock'


if not sys.stdin.isatty():
	pause = lambda: sys.stdin.read(1)
	catch_escape = False
else:
	try:
		from select import select
		import termios, tty, atexit
	except ImportError:
		pause = None
		catch_escape = False
	else:
		catch_escape = True
		def cleanup(old=termios.tcgetattr(sys.stdin.fileno())):
			termios.tcsetattr(sys.stdin.fileno(), termios.TCSAFLUSH, old)
		atexit.register(cleanup)
		tty.setcbreak(sys.stdin.fileno())
		def pause():
			sys.stdin.read(1)

try:
	from signal import SIGCHLD, SIG_DFL, signal, set_wakeup_fd
	from select import select, error as SelectError
	from errno import EAGAIN, EINTR
	from fcntl import fcntl, F_SETFD, F_GETFD, F_SETFL, F_GETFL
	from os import O_NONBLOCK
	try:
		from signal import siginterrupt
	except ImportError:
		# Sucks.
		siginterrupt = lambda signalnum, flag: None
	try:
		import cPickle as pickle
	except ImportError:
		import pickle
except ImportError:
	def call(*args, **kwargs):
		case = kwargs.pop('case')
		try:
			case.process = Popen(*args, **kwargs)
		except OSError:
			raise testcases.CannotStartTestee(sys.exc_info()[1])
		case.time_started = clock()
		if not case.maxwalltime:
			while True:
				exitcode, now = case.process.poll(), clock()
				if exitcode is not None:
					case.time_stopped = now
					break
				else:
					time.sleep(.001)
		else:
			time_end = case.time_started + case.maxwalltime
			while True:
				exitcode, now = case.process.poll(), clock()
				if exitcode is not None:
					case.time_stopped = now
					break
				elif now >= time_end:
					raise TimeLimitExceeded
				else:
					time.sleep(.001)
else:
	try:
		from fcntl import FD_CLOEXEC
	except ImportError:
		FD_CLOEXEC = 1
	
	try:
		from resource import getrusage, RUSAGE_SELF, RUSAGE_CHILDREN
	except ImportError:
		from time import clock as cpuclock
		getrusage = lambda who: None
	else:
		def cpuclock():
			rusage = getrusage(RUSAGE_SELF)
			return rusage.ru_utime + rusage.ru_stime
	
	try:
		from resource import setrlimit
		try:
			from resource import RLIMIT_AS
		except ImportError:
			from resource import RLIMIT_VMEM
	except ImportError:
		setrlimit = None
	
	# Make SIGCHLD interrupt sleep() and select()
	sigchld_pipe_read, sigchld_pipe_write = os.pipe()
	fcntl(sigchld_pipe_read, F_SETFL,
	      fcntl(sigchld_pipe_read, F_GETFL) | O_NONBLOCK)
	fcntl(sigchld_pipe_write, F_SETFL,
	      fcntl(sigchld_pipe_write, F_GETFL) | O_NONBLOCK)
	def bury_child(signum, frame):
		try:
			bury_child.case.time_stopped = clock()
		except Exception:
			pass
	signal(SIGCHLD, bury_child)
	set_wakeup_fd(sigchld_pipe_write)
	class SignalIgnorer(object):
		def __enter__(self):
			signal(SIGCHLD, SIG_DFL)
		def __exit__(self, exc_type, exc_value, traceback):
			signal(SIGCHLD, bury_child)
	signal_ignorer = SignalIgnorer()
	__all__ += 'signal_ignorer',
	
	# If you want this to work portably, don't set any stdio argument to PIPE
	def call(*args, **kwargs):
		global last_rusage
		bury_child.case = case = kwargs.pop('case')
		read, write = os.pipe()
		fcntl(write, F_SETFD, fcntl(write, F_GETFD) | FD_CLOEXEC)
		def preexec_fn():
			os.close(read)
			if setrlimit and case.maxmemory:
				maxmemory = ceil(case.maxmemory * 1048576)
				setrlimit(RLIMIT_AS, (maxmemory, maxmemory))
				# I would also set a CPU time limit but I do not want the time
				# passing between the calls to fork and exec to be counted in
			os.write(write, pickle.dumps((clock(), cpuclock()), 1))
		kwargs['preexec_fn'] = preexec_fn
		old_rusage = getrusage(RUSAGE_CHILDREN)
		last_rusage = None
		while True:
			try:
				os.read(sigchld_pipe_read, 512)
			except OSError:
				if sys.exc_info()[1].errno == EAGAIN:
					break
				else:
					raise
		siginterrupt(SIGCHLD, False)
		try:
			case.process = Popen(*args, **kwargs)
		except OSError:
			os.close(read)
			raise testcases.CannotStartTestee(sys.exc_info()[1])
		finally:
			siginterrupt(SIGCHLD, True)
			os.close(write)
		try:
			if not catch_escape:
				if case.maxwalltime:
					try:
						select((sigchld_pipe_read,), (), (), case.maxwalltime)
					except SelectError:
						if sys.exc_info()[1].args[0] != EINTR:
							raise
					if case.process.poll() is None:
						raise testcases.WallTimeLimitExceeded
				else:
					case.process.wait()
			else:
				if not case.maxwalltime:
					try:
						while case.process.poll() is None:
							s = select((sys.stdin, sigchld_pipe_read), (), ())
							if (sigchld_pipe_read not in s[0] and
							    sys.stdin.read(1) == '\33'):
								raise testcases.CanceledByUser
					except (SelectError, IOError):
						if sys.exc_info()[1].args[0] != EINTR:
							raise
						else:
							case.process.poll()
				else:
					time_end = clock() + case.maxwalltime
					try:
						while case.process.poll() is None:
							remaining = time_end - clock()
							if remaining > 0:
								s = select((sys.stdin, sigchld_pipe_read),
								           (), (), remaining)
								if (sigchld_pipe_read not in s[0] and
								    sys.stdin.read(1) == '\33'):
									raise testcases.CanceledByUser
							else:
								raise testcases.WallTimeLimitExceeded
					except (SelectError, IOError):
						if sys.exc_info()[1].args[0] != EINTR:
							raise
						else:
							case.process.poll()
		finally:
			case.time_started, cpustart = pickle.loads(os.read(read, 512))
			os.close(read)
			del bury_child.case
		new_rusage = getrusage(RUSAGE_CHILDREN)
		if (case.maxwalltime and
		    case.time_stopped - case.time_started > case.maxwalltime):
			raise testcases.WallTimeLimitExceeded
		if new_rusage:
			time_started = old_rusage.ru_utime + old_rusage.ru_stime + cpustart
			time_stopped = new_rusage.ru_utime + new_rusage.ru_stime
			if case.maxcputime or not case.maxwalltime:
				case.time_started = time_started
				case.time_stopped = time_stopped
				case.time_limit_string = case.cpu_time_limit_string
				if (case.maxcputime and
				    time_stopped - time_started > case.maxcputime):
					raise testcases.CPUTimeLimitExceeded
		if case.maxmemory:
			if sys.platform != 'darwin':
				maxrss = case.maxmemory * 1024
			else:
				maxrss = case.maxmemory * 1048576
			if last_rusage and last_rusage.ru_maxrss > maxrss:
				raise testcases.MemoryLimitExceeded
			elif (new_rusage and
			      new_rusage.ru_maxrss > old_rusage.ru_maxrss and
			      new_rusage.ru_maxrss > maxrss):
				raise testcases.MemoryLimitExceeded

# Emulate memory limits on platforms compatible with 4.3BSD but not XSI
# I say 'emulate' because the OS will allow excessive memory usage
# anyway; Upreckon will just treat the test case as not passed.
# To do this, we not only require os.wait4 to be present but also
# assume things about the implementation of subprocess.Popen.
try:
	def waitpid_emu(pid, options, _wait4=os.wait4):
		global last_rusage
		pid, status, last_rusage = _wait4(pid, options)
		return pid, status
	_waitpid = os.waitpid
	os.waitpid = waitpid_emu
	try:
		defaults = Popen._internal_poll.__func__.__defaults__
	except AttributeError:
		# Python 2.5
		defaults = Popen._internal_poll.im_func.func_defaults
	i = defaults.index(_waitpid)
	defaults = defaults[:i] + (waitpid_emu,) + defaults[i+1:]
	try:
		Popen._internal_poll.__func__.__defaults__ = defaults
	except AttributeError:
		# Python 2.5 again
		Popen._internal_poll.im_func.func_defaults = defaults
except (AttributeError, ValueError):
	pass


def kill(process):
	try:
		process.kill()
	except AttributeError:
		os.kill(process.pid, SIGKILL)


def terminate(process):
	try:
		process.terminate()
	except AttributeError:
		os.kill(process.pid, SIGTERM)