
__author__ = 'Miroslav Bauer'


import zmq
import util
import pickle
import simplejson
from logger     import Logger
from objproxy   import RpcProxy
from zmqthread  import ZmqThread
from requestapi import RequestAPI
from resman     import ResourceManager
from constants  import JOB_EXTERNAL_STATE

log = Logger('dmorph.rhandler.work')


class RequestWorker(ZmqThread, RequestAPI):
	''' This class fetches and executes user requests
	    (RequestAPI method invocations)
	'''

        def __init__(self, context):
		super(RequestWorker, self).__init__(context)
		self.name   = 'requestworker-%s' % self.name.replace('Thread-','')

        def _connect(self):
		super(RequestWorker, self)._connect()
		self._sockets['fetcher']   = self._context.socket(zmq.PULL)
		self._sockets['sender']    = self._context.socket(zmq.PUSH)
		self._sockets['scheduler'] = self._context.socket(zmq.DEALER)
		self._sockets['scheduler'].identity = self.name
		self.resman = RpcProxy(self._context, ResourceManager)
			
		try:
			self._sockets['fetcher'].connect('inproc:///requests/workers/tasks')
			self._sockets['sender'].connect('inproc:///requests/workers/responses')
			self._sockets['scheduler'].connect('inproc:///scheduler/responses')
			self.resman.connect('resman')
		except IOError:
			log.err('Could not connect to control sockets,'\
				'RequestHandler must be started first')
			self.terminate()
			raise

	def _sendResponse(self, recipient, message):
		self._sockets['sender'].send_multipart(util.flatten([recipient, message]))

        def run(self):
		''' Fetch method invocation data from request and execute method '''
                self._connect()
		self._signalReady()

                while not self.stop.is_set():
                        originator, cmd, args = self._sockets['fetcher'].recv_multipart()
			args = pickle.loads(args)

                        method = getattr(self, cmd, None) if hasattr(RequestAPI, cmd) else None
                        if (method == None):
                                log.err('Client %s tried to invoke unknown method '\
					'"%s"' % (originator, cmd))
				self.error(originator, 'TypeError', 'Invalid command "%s"' % cmd)
                                continue
			log.debug('Running %s on %s' % (method.__name__, self.name))
			try:
				if args == []:
					result = method(originator)
				else:
					result = method(originator, *args)
			except Exception, e:
				result = self.error(originator, e.__class__.__name__, e.message)

			if result != None:	
				response = util.flatten([result])
				self._sendResponse(originator, response)

			log.debug('%s is done with %s' % (self.name, cmd))

		self._cleanup()

        def compute(self, caller, jobId, jd):
                if not jobId.isdigit():
                        return 'Invalid job ID %s' % jobId
                jid = int(jobId)
		jobDct = simplejson.loads(jd)
		job = util.toJob(jobDct)
		job.setId(jid)
		self._sockets['scheduler'].send(pickle.dumps(job))

		result = None
		notice = None
		self._subscribeFor([str(job.jid)])
		scheduler = self._sockets['scheduler']
		notifs    = self._sockets['notifications_sub']
		pl = util.createPoller([scheduler, notifs])

		while True:
			socks = dict(pl.poll())
			if scheduler in socks and socks[scheduler] == zmq.POLLIN:
				result = pickle.loads(scheduler.recv())
				if isinstance(result, Exception):
					raise result
				else:
					return self.results(caller, jid, result)

			elif notifs in socks and socks[notifs] == zmq.POLLIN:
				notice = self._recvNotification()
				if notice.startswith('STATE'):
					code = int(notice.split(':')[1])
					n = self.notify(caller, jid,
						util.stateToStr(code, JOB_EXTERNAL_STATE))
					self._sendResponse(caller, n)
		self._disconnect()

	def reloadagents(self, caller):
		self.resman.reloadResources()
		
	def getagents(self, caller):
		return ['AGENTS', simplejson.dumps(self.resman.getResources())]

	def error(self, caller, eType, eMsg):
		self._sendResponse(caller, super(RequestWorker, self).error(caller, eType, eMsg))
	
	def terminate(self):
		super(RequestWorker, self).terminate()
		self.resman._terminate()
