# vim: set nolist noet ts=4:
# Copyright 2002, 2003, 2004, 2005 Ben Escoto
#
# This file is part of rdiff-backup.
#
# rdiff-backup is free software; you can redistribute it and/or modify
# under the terms of the GNU General Public License as published by the
# Free Software Foundation; either version 2 of the License, or (at your
# option) any later version.
#
# rdiff-backup is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with rdiff-backup; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307
# USA
from collections import namedtuple
from multiprocessing import cpu_count
import Queue
import itertools
import threading

Job = namedtuple('Job', ['func', 'iterable', 'outqueue'])
Task = namedtuple('Task', ['func', 'args', 'index', 'outqueue'])
Result = namedtuple('Result', ['index', 'value'])

RUNNING = 0
STOPPED = 1


def worker(taskqueue):
	while True:
		task = taskqueue.get(True)
		if task is None:
			taskqueue.task_done()
			break

		if task.func is None:
			# It means this task was the last of an iterable job
			result = None
		else:
			value = task.func(*task.args)
			result = Result(task.index, value)

		task.outqueue.put(result, block=True)

		taskqueue.task_done()


class Pool(object):
	def __init__(self, processes=None,
			     max_taskqueue_size=0, max_jobqueue_size=0):
		if processes is None:
			processes = cpu_count()
		self.processes = processes

		self.state = STOPPED

		# Init queues
		self.taskqueue = Queue.Queue(maxsize=max_taskqueue_size)
		self.jobqueue = Queue.Queue(maxsize=max_jobqueue_size)

		# Init workers
		self.workers = []
		self.start_workers()

		# Init task handler thread
		self._job_handler_thread = self._start_handler_thread(self._job_handler,
															   self.jobqueue,
															   self.taskqueue)

	def start_workers(self):
		while len(self.workers) < self.processes:
			w = self._start_handler_thread(worker, self.taskqueue)
			self.workers.append(w)

		for w in self.workers:
			if not w.is_alive():
				w.start()

	def _start_handler_thread(self, func, *args):
		thread = threading.Thread(target=func, args=args)
		thread.daemon = True
		thread.start()

		return thread

	def create_job(self, func, iterable, max_outqueue_size=0):
		outqueue = Queue.Queue(maxsize=max_outqueue_size)

		job = Job(func, iterable, outqueue)
		self.jobqueue.put(job)

		return job


	def imap(self, func, iterable, max_outqueue_size=0):
		iterable = itertools.imap(None, iterable)

		job = self.create_job(func, iterable,
							  max_outqueue_size=max_outqueue_size)

		return IMapIterator(job.outqueue)

	def apply(self, func, *args):
		job = self.create_job(func, [args])

		return AsyncResult(job.outqueue)

	def stop(self):
		self.jobqueue.put(None, block=True)

	def join(self, timeout=None):
		self.stop()
		self._job_handler_thread.join(timeout=timeout)
		for w in self.workers:
			w.join(timeout=timeout)

	def _job_handler(self, jobqueue, taskqueue):
		while True:
			job = jobqueue.get(True)
			if job is None:
				for w in self.workers:
					taskqueue.put(None)
				break

			for (index, args) in enumerate(job.iterable):
				task = Task(job.func, args, index, job.outqueue)
				taskqueue.put(task, block=True)
			taskqueue.put(Task(None, None, None, job.outqueue), block=True)

			jobqueue.task_done()

class IMapIterator(object):
	def __init__(self, outqueue):
		self.outqueue = outqueue
		self.results = {}
		self.index = 0

	def __iter__(self):
		return self

	def next(self):
		while True:
			if self.index in self.results:
				result = self.results.pop(self.index)
			else:
				result = self.outqueue.get(True)
				if result is None:
					raise StopIteration()

			if result.index != self.index:
				self.results[result.index] = result
				continue

			self.index += 1
			return result.value

class AsyncResult(object):
	def __init__(self, outqueue):
		self.outqueue = outqueue
		self.completed = False
		self.value = None

	def wait(self):
		if self.completed:
			return
		self.value = self.outqueue.get(True)
		self.completed = True

	def get(self):
		self.wait()
		return self.value
