"""This is designed to let you do asynchronous I/O conveniently.
"""

from __future__ import with_statement

import sys
import threading
import die
import system_load as SL
from gyropy import g_mailbox

Lock = threading.Lock

class CaughtException(object):
	def __init__(self, x):
		self.ex_type, self.ex_value, self.ex_traceback = x

	def __repr__(self):
		return "CaughtException: %s" % str(self.ex_value)

	def reraise(self):
		raise self.ex_type, self.ex_value, self.ex_traceback
	
	get = reraise


class threading_with_result(threading.Thread):
	def __init__(self, tag=None, group=None, target=None, name=None, args=(), kwargs={}):
		threading.Thread.__init__(self, group=group, target=target, name=name, args=args, kwargs=kwargs)
		self.rv = None
		self.rvset = False
		self.tag = tag

	def run(self):
		try:
			self.rv = self._Thread__target(*self._Thread__args, **self._Thread__kwargs)
		except:
			self.rv = CaughtException(sys.exc_info())
		self.rvset = True

	def get(self):
		assert self.rvset
		check_n_raise(self.rv)
		return self.rv


class Thread_pool(object):
	EOF = g_mailbox.EOF
	CaughtException = CaughtException

	def __init__(self, n=None):
		if n is None:
			n = SL.ncpu()
		assert n > 0
		self.iqueue = g_mailbox.listbox()
		self.t = []
		for i in range(n):
			tt = threading.Thread(target=self._target, args=(self.iqueue,) )
			tt.start()
			self.t.append(tt)
	
	def __len__(self):
		return len(self.t)

	def startRV(self, target=None, args=None, kwargs=None):
		assert target is not None
		wbox = g_mailbox.waitbox()
		self.iqueue.put((wbox, target, args, kwargs))
		return wbox

	def startNR(self, target=None, args=None, kwargs=None):
		assert target is not None
		self.iqueue.put((None, target, args, kwargs))
	
	def startQ(self, queue, tag, target=None, args=None, kwargs=None):
		def wrapped_target(target, *args, **kwargs):
			return (tag, target(*args, **kwargs))
		self.iqueue.put((queue, wrapped_target, args, kwargs))

	@classmethod
	def _target(cls, iqueue):
		while True:
			try:
				wbox, target, args, kwargs = iqueue.get()
			except cls.EOF:
				break
			if kwargs is None:
				kwargs = {}
			if args is None:
				args = ()

			try:
				if wbox is not None:
					wbox.put(target(*args, **kwargs))
				else:
					target(*args, **kwargs)
			except:
				if wbox is not None:
					wbox.put( cls.CaughtException(sys.exc_info()) )
				else:
					die.catch("Uncaught exception in Thread_pool thread")
	
	def join(self):
		if self.t:
			self.iqueue.close()
			while self.t:
				tmp = self.t.pop(0)
				tmp.join()
	
	def __del__(self):
		self.join()
		



def to_be_joined(tlist):
	"""Pick a ripe thread for joining.
	@type tlist: list(L{threading.Thread})
	@param tlist: a list of threads.
	@return: a thread that is either finished and ready to be joined,
		or (if none are ready yet) the oldest thread.
	@rtype: L{threading.Thread}
	"""
	for (i, t) in enumerate(tlist):
		if not t.isAlive():
			return tlist.pop(i)
	return tlist.pop(0)


class Thread_poolW(Thread_pool):
	"""Here, you have a pool of n threads.   You can call
	C{x.start(function, args, kwargs)} to start something running
	(perhaps a write to a file), and then go off and compute
	something else while the I/O completes.
	If you need to wait for completion, call C{x.join()}.

	This class is not designed to return values from the function.
	"""

	def __init__(self, n=None):
		Thread_pool.__init__(self, n)


	def start(self, fn, args=None, kwargs=None):
		self.startNR(fn, args, kwargs)




Thread_poolNR = Thread_poolW	# Obsolete


def check_n_raise(x):
	if isinstance(x, CaughtException):
		x.reraise()
	return x


class Thread_poolR(Thread_pool):
	"""Here, you have a pool of n threads.   You can call
	C{x.start(function, args, kwargs)} to start something running
	(perhaps a write to a file), and then go off and compute
	something else while the I/O completes.
	If you need to wait for completion, call C{x.join()}.

	This class is designed to return values from the function via
	C{get()} or C{getany()}.
	"""

	def __init__(self, n=None):
		Thread_pool.__init__(self, n)
		self.lock = threading.Lock()
		self.oqueue = g_mailbox.mailbox()
		self.tags = {}


	def start(self, tag, fn, args=None, kwargs=None):
		with self.lock:
			self.tags[tag] = self.startRV(target=fn, args=args, kwargs=kwargs)


	def getany(self):
		"""
		@return: the first available answer, as a (tag, value) pair.   The tag was set
			in the call to C{start} and the value is the result of the computation.
			If the computation raised an exectption, C{getany} will return a CaughtException
			object instead.
		@rtype: a tuple(whatever, whatever) pair from a computation or a L{CaughtException}
		"""
		with self.lock:
			for tag in self.tags.keys():
				if self.tags[tag]:
					return (tag, self.tags.pop(tag).get())
		# This is not quite right.   It picks a waitbox to wait on, but it
		# really should defer that decision until one has an answer available.
		tag, v = self.tags.popitem()
		return (tag, v.get())


	def get(self, tag):
		"""
		@rtype: whatever or a L{CaughtException}
		"""
		with self.lock:
			return self.tags.pop(tag).get()


	def has_answer(self):
		"""Is there an answer ready?"""
		with self.lock:
			for (k, v) in self.tags.items():
				if len(v):
					return True
		return False


	# def loaded(self):
		# """Have we loaded up all the processors?"""
		# with self.lock:
			# # This is not quite right.  It only returns True when there
			# # is one more job than processor.
			# return len(self.iqueue) > 0



def pairmap(fn, inlist, *args, **kwargs):
	tp = Thread_poolR(kwargs.get('_poolsize', None))
	no = 0
	s = 0
	for (j,xi) in enumerate(inlist):
		tp.start(xi, fn, (xi,)+args, kwargs)
		s += 1
		if s > len(tp):
			yield tp.getany()
			no += 1
	while no < s:
		yield tp.getany()
		no += 1
	

def testmap():
	q = list(pairmap(lambda x:x, range(10)))
	assert len(q)==10
	seen = set()
	for (i,j) in q:
		assert i==j
		seen.add(i)
	assert sorted(seen)==range(10)


if __name__ == '__main__':
	testmap()
