
__author__ = 'Miroslav Bauer'


import zmq
import pickle
import time
from logger import Logger
from threading import Thread, Event
from constants import WAITING_TIMEOUT

log = Logger('dmorph.thread')

class ZmqThread(Thread):
	''' Threaded class with ZMQ communication capabilities '''

	def __init__(self, context, parent=None, hasChildren=False):
		super(ZmqThread, self).__init__()
		self.stop = Event()
		self.setDaemon(True)
		self._context    = context
		self._sockets    = {}
		self._threads    = {}
		self.parent      = parent
		self.hasChildren = self._threads != {} or hasChildren

	def _bind(self):
		if self.hasChildren:
			self._sockets['children'] = self._context.socket(zmq.PULL)
			try:
				self._sockets['children'].bind('inproc:///parent/%s' % self.name)
			except IOError:
				self.terminate()
				raise

	def _connect(self):
		if self.parent:
			self._sockets['parent'] = self._context.socket(zmq.PUSH)
		self._sockets['notifications_pub'] = self._context.socket(zmq.PUB)
		self._sockets['notifications_sub'] = self._context.socket(zmq.SUB)
		try:
			if self.parent:
				self._sockets['parent'].connect('inproc:///parent/%s' % self.parent)
			self._sockets['notifications_pub'].connect('inproc:///notifications_in')
		except IOError:
			self.terminate()
			raise

	def getName(self):
		return self.name

	def notice(self, topic, message, toParent=False):
		''' Publishes a message with topic via zmq.PUB socket '''
		if toParent:
			self._sockets['parent'].send('%s %s' % (topic, message))
		else:
			self._sockets['notifications_pub'].send('%s %s' % (topic, message))
	
	def _disconnect(self):
		self._sockets['notifications_sub'].disconnect('inproc:///notifications_out')

	def _subscribeFor(self, topics=[]):
		try:
			self._disconnect()
		except zmq.ZMQError:
			pass
		for topic in topics:
			self._sockets['notifications_sub'].setsockopt(zmq.SUBSCRIBE, topic)
		self._sockets['notifications_sub'].connect('inproc:///notifications_out')
	
	def _recvNotification(self, fromChilds=False):
		return self._recvRawNotification(fromChilds)[1]
	
	def _recvRawNotification(self, fromChilds=False):
		if fromChilds:
			return self._sockets['children'].recv().split(' ', 1)
		else:
			return  self._sockets['notifications_sub'].recv().split(' ', 1)

	def _signalReady(self, toParent=False):
		self.notice(self.name, 'READY', toParent)

	def _waitFor(self, readySenders=[], forChildren=False):
		starttime = time.time()
		self._subscribeFor(readySenders)
		while ((time.time()-WAITING_TIMEOUT) < starttime) and readySenders:
			log.debug('%s is waiting for %s' % (self.name, readySenders))
			sender, message = self._recvRawNotification(forChildren)
			if message == 'READY' and sender in readySenders:
				log.debug('%s got READY from %s' % (self.name, sender))
				readySenders.remove(sender)

	def _cleanup(self):
		[thr.terminate() for tname, thr in self._threads.items()]
		[skt.close() for sname, skt in self._sockets.items()]

	def terminate(self):
		''' Closes all sockets and children threads '''
		self.stop.set()
		self._cleanup()
		log.info('Thread %s has been terminated!' % self.name)

class RpcZmqThread(ZmqThread):
	''' ZmqThread class capable of handling Remote Procedure Calls
	    from RpcProxy objects.
	'''
	
	def __init__(self, context, parent=None, hasChildren=False):
		super(RpcZmqThread, self).__init__(context, parent, hasChildren)

	def _bind(self, endpoint):
		super(RpcZmqThread, self)._bind()
		self._sockets['responder'] = self._context.socket(zmq.REP)
                try:
                        self._sockets['responder'].bind('inproc://%s' % endpoint)
                except IOError:
                        self.terminate()
                        raise

	def _handleMethodInvocation(self, async=False):
		call = self._sockets['responder'].recv_multipart()
		if async:
			self._sockets['responder'].send(None)
			return
		try:
			args = None
			if len(call) == 2:
				cmd, args = call
				args = pickle.loads(args)
			elif len(call) == 1:
				cmd = call[0]
			else:
				raise TypeError('Invalid call: %s' % call)
			if cmd.startswith('_'):
				raise TypeError('Tried to invoke private method: %s' % cmd)
			method = getattr(self, cmd, None) if hasattr(self, cmd) else None
			if (method == None):
				raise RuntimeError('Tried to invoke unknown method "%s"' % cmd)
			if args:
				result = method(*args)
			else:
				result = method()
		except Exception, e:	
			result = e
			log.err('%s:%s' % (e.__class__.__name__, e.message))

		self._sockets['responder'].send(pickle.dumps(result))
