
__author__ = 'Miroslav Bauer'


import os
import zmq
import util
import pickle
import threading
from logger    import Logger
from rworker   import RequestWorker
from constants import (
	CTRL_PROTO, CTRL_ENDPOINT, ALLOWED_IDS, NUM_WORKERS,
	FILES_OWNER_UID, FILES_OWNER_GID
)

log = Logger('dmorph.rhandler')

class RequestHandler(threading.Thread):
	''' Application's front-end class for accepting user
	    input (RequestAPI method invocations) and for
	    reporting back to user
	'''

	def __init__(self, context):
		super(RequestHandler, self).__init__()
		self.setDaemon(True)
		self.name      = 'rhandler'
		self._context  = context
		self._workers = []

	def _bind(self):
		try:
			self._responder = self._context.socket(zmq.ROUTER)
			self._distributor = self._context.socket(zmq.PUSH)  
			self._collector = self._context.socket(zmq.PULL)

			self._responder.bind("%s://%s" % (CTRL_PROTO, CTRL_ENDPOINT))
			self._responder.bind('inproc:///requests/responses')
			self._distributor.bind('inproc:///requests/workers/tasks')
			self._collector.bind('inproc:///requests/workers/responses')

			os.chown(CTRL_ENDPOINT, FILES_OWNER_UID, FILES_OWNER_GID)

		except IOError:
			log.err('Could not create control sockets')
			self.terminate()
			raise

	def run(self):
		''' Starts request worker threads and starts polling for requests,
		    responses and distributing tasks to workers
		'''
		self._bind()
		self._workers = [RequestWorker(self._context) for i in range(NUM_WORKERS)]
		for worker in self._workers:
			worker.start() 

		poll = util.createPoller([self._responder, self._collector])

		log.info('Request handler started!')
		log.debug('Using %s request worker threads' % NUM_WORKERS)
		log.info('Waiting for requests on: %s://%s' % (CTRL_PROTO, CTRL_ENDPOINT))

		while True:
			socks = dict(poll.poll())
			if self._responder in socks and socks[self._responder] == zmq.POLLIN:
				self._handleRequest(self._responder.recv_multipart())
			elif self._collector in socks and socks[self._collector] == zmq.POLLIN:
				data = self._collector.recv_multipart()
				log.debug('Sending "%s" to: %s' % (data[1:], data[0]))
				self._responder.send_multipart(data)

	def _handleRequest(self, request):
		if len(request) < 2:
			log.err('Invalid request: %s' % request)
			return

		originator, cmd = request[:2]
		cmd = cmd.lower()
		if not originator in ALLOWED_IDS:
			log.err('Identity %s not in ALLOWED_IDS' % originator)
			return

		self._distributor.send_multipart([originator, cmd, pickle.dumps(request[2:])])

	def terminate(self):
		''' Shut down all sockets and threads '''
		if hasattr(self, 'self._responder'): self._responder.close()
		if hasattr(self, 'self._distributor'): self._distributor.close()
		if hasattr(self, 'self._collector'): self._collector.close()
		[worker.terminate() for worker in self._workers]
