
__author__ = 'Miroslav Bauer'


import os
import zmq
import saga
import time
import pickle
from logger    import Logger
from os.path   import basename, dirname
from objproxy  import RpcProxy
from zmqthread import ZmqThread
from resource  import ResourceAgent
from constants import JOB_EXTERNAL_STATE as state
from constants import FILES_OWNER_UID, FILES_OWNER_GID, CACHING_ENABLED

log = Logger('dmorph.job.executor')

class JobExecutor(ZmqThread):
	''' Thread responsible for the whole job preparation, execution,
	    monitoring and collecting results process on the resource
	    given.
	'''

	def __init__(self, context, job, resourcename):
		super(JobExecutor, self).__init__(context)
		self.job          = job
		self.resourcename = resourcename
		self.name         = 'exec-job%s-res-%s' % (job.jid, resourcename)
		self.service      = None
	
	def _connect(self):
		super(JobExecutor, self)._connect()
		self._sockets['sender'] = self._context.socket(zmq.PUSH)
		try:
			self._sockets['sender'].connect('inproc:///scheduler/executors')
			self.resource.connect('resource/%s' % self.resourcename)
		except IOError:
			log.err('Could not connect job executor sockets')
			self.terminate()
			raise

	def _getFromCache(self, localfiles):
		return self.resource.getFromCache(localfiles)

	def _changeJobState(self, st):
		if st == state['preparing']:
			self.resource.incStat('preparing_jobs')
		elif st == state['executing']:
			self.resource.decStat('preparing_jobs')
			self.resource.incStat('executing_jobs')
		elif st == state['finishing']:
			self.resource.decStat('executing_jobs')
			self.resource.incStat('unstaging_jobs')
		elif st == state['finished']:
			self.resource.decStat('unstaging_jobs')
			self.resource.incStat('jobs_processed')
		self.job.state = st
		self.notice(str(self.job.jid), 'STATE:%s' % st)

	def _sendResult(self, response):
		self._sockets['sender'].send(pickle.dumps([self.job.jid, response]))

	def _initSagaContext(self):
		method, data = self.resource.getSagaAuthData()
		if not method: return
		ctx = saga.Context(method)
		for key, value in data.items():
			setattr(ctx, key, value)
		self.session.add_context(ctx)

	def _initResourceService(self):
		url = self.resource.getServiceUrl()
		self.service = saga.job.Service(url, session=self.session)

	def _prepareSagaJob(self):
		self._changeJobState(state['preparing'])
		self._initSagaContext()
		self._initResourceService()
		jd = self.job.getSagaDescription(self.resource.getAttrs())
		return (jd, self.service.create_job(jd))

	def run(self):
		''' Submits job to resource and collects computation results '''
		ts_start      = time.time()
		ts_xfer_start = 0
		xfer_time     = 0

		self.session  = saga.Session()
		self.resource = RpcProxy(self._context, ResourceAgent)
		self._connect()
		try:
			jDescription, sagaJob = self._prepareSagaJob()
			endpoint    = self.resource.getServiceEndpoint()
			xfer        = self.resource.getXferBackend()
			serviceHost = '%s://%s' % (xfer, endpoint)
			jobBaseDir  = '%s%s' % (serviceHost, jDescription.working_directory)
			if CACHING_ENABLED:
				# Check for files in resource's cache
				cacheDir = self.resource.getCacheDir()
				ts_xfer_start = time.time()
				self._stagein(jobBaseDir, cacheDir)
			else:
				ts_xfer_start = time.time()
				self._stagein(jobBaseDir)
			ts_xfer_end = time.time()

			sagaJob.run()
			self._changeJobState(state['executing'])
			log.info('Running #%s as SID #%s' % (self.job.jid, sagaJob.id))

			sagaJob.wait()
			self._changeJobState(state['finishing'])
			
			results = {}
			# Collect job results from stdout, stder and output files
			for pos, par in  self.job.params.items():
				if par.direction == 'in': continue
				val = None
				if pos == 'stdout':
					val = self._getJobOutput(jDescription, serviceHost)
				elif pos == 'stderr':
					val = self._getJobError(jDescription, serviceHost)
				else:
					val = par.value
				
				results[par.name] = val
			self._stageout(jobBaseDir)

			log.info('Job SID #%s finished!' % sagaJob.id)
			log.info('Job SID #%s Exitcode: %s' % (sagaJob.id, sagaJob.exit_code))
			self._changeJobState(state['finished'])

			if ts_xfer_start:
				xfer_time = time.time() - ts_xfer_start
				self.resource.updateStat('transfer_speed_avg', self.job.size/xfer_time)

			self._sendResult(results)
		except Exception, e:
			self._sendResult(e)
			self.resource.incStat('jobs_failed')
			self.terminate()
			raise

		exec_time = time.time() - ts_start
		self.resource.updateStat('exec_time_avg', exec_time - xfer_time)
		self.terminate()
	
	def _stagein(self, target, cachedir=None):
		log.info('Staging files to: %s' % target)
	
		log.info('Creating job #%s working directory' % self.job.jid)
		saga.filesystem.Directory('%s/..' % target).make_dir(target)
		srcs = self._getAllSourceFiles(self.job.infiles)
		copylist = ['file://%s' % f for f in srcs]	

		if CACHING_ENABLED:
			csrcs = self._getFromCache(srcs)
			for c, filepath in csrcs:
				# File has been just cached, copy it to cache dir
				if c: self._copyRecursive(['file://%s' % filepath[1]], filepath[0])
				# Copy file from cache to job's directory
				self._copyRecursive([filepath[0]], target+basename(filepath[1]))
		else:	
			self._copyRecursive(copylist, target)

	def _getAllSourceFiles(self, flist, prefix=''):
		allfiles = []
		for f in flist:
			f = prefix + f
			if os.path.isdir(f):
				dirfiles = os.listdir(f)
				dirflist = self._getAllSourceFiles(dirfiles, f)
				allfiles += dirflist
			else:
				allfiles.append(f)
		return allfiles

	def _copyRecursive(self, flist, target, prefix='', direction='in'):
		for fn in flist:
			fn = '%s%s' % (prefix, fn)
			f = saga.filesystem.File(fn, session=self.session)
			if f.is_file():
				log.debug('Staging %s: %s' % (direction, fn))
				f.copy(target) if direction == 'in' else f.move(target)
			else:
				d = saga.filesystem.Directory(fn, session=self.session)
				if direction == 'in':
					self._copyRecursive(d.list(), target, fn, direction)
				else:
					d.move(target)
	
	def _stageout(self, source):
		for fn in self.job.outfiles:
			source = '%s%s' % (source, basename(fn))
			target = 'file://%s' % fn
			self._copyRecursive([source], target, direction='out')
			os.chown(fn, FILES_OWNER_UID, FILES_OWNER_GID)
		
		log.info('Removing job #%s working directory' % self.job.jid)
		saga.filesystem.Directory(dirname(source)).remove(flags=saga.filesystem.RECURSIVE)
	
	def _getJobOutput(self, jobDesc, hostRoot):
		out = saga.filesystem.File('%s%s' % (hostRoot, jobDesc.output))
		return out.read()

	def _getJobError(self, jobDesc, hostRoot):
		out = saga.filesystem.File('%s%s' % (hostRoot, jobDesc.error))
		return out.read()

	def terminate(self):
		''' Sets stop signal and disconnects resource proxy object '''
		self.resource._terminate()
		super(JobExecutor, self).terminate()
