# 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
import sync

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


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

		if callable(task.out):
			args = task.options.get('out_args', ())
			out = task.out(*args)
		else:
			out = task.out

		try:
			value = task.func(*task.args)
		except Exception as e:
			value = e
		result = value

		out.set(result)
		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

		# 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)

	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, **options):
		max_outqueue_size = options.pop('max_outqueue_size', 0)
		outqueue = JobResultQueue(maxsize=max_outqueue_size)

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

		return job

	def create_task(self, func, args, out=None, **options):
		if out is None:
			out = sync.AsyncValue()

		task = Task(func, args, out, options)
		self.taskqueue.put(task, block=True)

		return task

	def imap(self, func, iterable, **options):
		iterable = itertools.imap(None, iterable)

		job = self.create_job(func, iterable, **options)

		return IMapIterator(job.outqueue)

	def apply(self, func, *args, **options):
		task = self.create_task(func, args, **options)

		return task.out

	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):
		while True:
			job = self.jobqueue.get(True)
			if job is None:
				for w in self.workers:
					self.taskqueue.put(None)
				break

			for (index, args) in enumerate(job.iterable):
				out = job.outqueue.slot(index)
				task = self.create_task(job.func, args, out, **job.options)
				#  task = self.create_task(job.func, args, job.outqueue.slot,
				#                          out_args=(index,), **job.options)
			job.outqueue.set_length(index + 1)

			self.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:
			slot = self.outqueue.get(self.index)
			if slot is None:
				raise StopIteration()

			result = slot.get()
			if isinstance(result, Exception):
				raise result

			self.index += 1
			return result

class JobResultQueue(object):
	def __init__(self, maxsize=None):
		self.maxsize = maxsize

		self.slots = {}
		self.updated = threading.Condition(threading.Lock())
		self.length = None

	def slot(self, index):
		self.updated.acquire()
		try:
			while len(self.slots) == self.maxsize:
				self.updated.wait()

			slot = self.slots[index] = sync.AsyncValue()
			self.updated.notify()
			return slot
		finally:
			self.updated.release()

	def get(self, index):
		self.updated.acquire()
		try:
			while index not in self.slots:
				if self.length is not None and index >= self.length:
					return None
				self.updated.wait()

			slot = self.slots.pop(index)
			self.updated.notify()
			return slot
		finally:
			self.updated.release()

	def set_length(self, length):
		self.updated.acquire()
		try:
			self.length = length
			self.updated.notify()
		finally:
			self.updated.release()
