
__author__ = 'Miroslav Bauer'


import os
import saga
import time
import pickle
import subprocess
from logger    import Logger
from os.path   import dirname
from filecache import FileCache
from zmqthread import RpcZmqThread
from constants import (
	RES_STATE, RES_STAT_WEIGHT,
	CACHING_ENABLED, RES_SAVE_PATH
)

log = Logger('dmorph.resource')

class ResourceAgent(RpcZmqThread):
	''' This class represents computational resource
	    with it's specific configuration.
	'''
	
	def __init__(self, context, name, enabled, endpoint, backend,
		auth, xfer, env={}, scratch=None, features={}):
		super(ResourceAgent, self).__init__(context)
		self.name        = name
		self.endpoint    = endpoint
		self.jobbackend  = backend
		self.auth        = auth
		self.environment = env
		self.scratchdir  = scratch
		self.features    = features
		self.xfer        = xfer

		self.filecache = None
		self.stats     = { stat : 0 for stat in RES_STAT_WEIGHT.keys() }
		self.stats['last_job_assignment_time'] = time.time()

	def run(self):
		''' Connects and binds it's sockets, loads its stats
		    and cache and starts handling RPC calls.
		'''
		self._connect()
		self._bind('resource/%s' % self.name)
		self._loadStats()

		if CACHING_ENABLED and not self.filecache:
			self.filecache = FileCache(xfer,
				'%s%s/filecache' % (self.endpoint, self.scratchdir))
		self._signalReady(toParent=True)

		log.info('Agent "%s" started!' % self.name)

		while not self.stop.is_set():
			self._handleMethodInvocation()

		self.terminate()

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

	def _saveStats(self):
		if not os.path.exists(dirname(RES_SAVE_PATH)):
			os.makedirs(dirname(RES_SAVE_PATH))
		with open('%s%s' % (RES_SAVE_PATH, self.name),'wb') as f:
			log.info('Saving %s stats to file' % self.name)
			if CACHING_ENABLED:
				pickle.dump([self.stats, self.filecache], f)
			else:
				pickle.dump([self.stats], f)

	def _loadStats(self):
		if not os.path.isfile('%s%s' % (RES_SAVE_PATH, self.name)):
			return
		with open('%s%s' % (RES_SAVE_PATH, self.name),'rb') as f:
			log.info('Loading %s stats from file' % self.name)
			savedList = pickle.load(f)
			self.stats = savedList[0]
			if CACHING_ENABLED and len(savedList) == 2:
				self.filecache = savedList[1]

	def _checkKrbTicket(self):
		return True if subprocess.call(['klist', '-s']) == 0 else False

	def incStat(self, statname):
		self.stats[statname] += 1
	
	def decStat(self, statname):
		self.stats[statname] -= 1
	
	def updateStat(self, statname, value):
		''' Updates metric to new value, or calculate new average value '''
		if statname.endswith('_avg'):
			oldavg = self.stats[statname]
			newavg = (oldavg + ((value - oldavg)/(self.stats['jobs_processed'])))
			self.stats[statname] = newavg
			log.debug('setting %s: avg:%s val:%s newavg:%s' % (statname,oldavg,value, newavg))
		else:	
			self.stats[statname] = value
			log.debug('setting %s=%s' % (statname,value))

	def isAlive(self):
		''' Checks if resource responds to ping '''
		session = saga.Session()
		pingcmd = ['ping', '-c', '1', self.endpoint.split(':')[0]]
		return True if subprocess.call(pingcmd, stdout=subprocess.PIPE) == 0 else False

	def isAccessible(self):
		''' Checks if it's possible to log in to resource with SSH '''
		try:
			homedir = saga.filesystem.Directory("%s://%s/"
				% (self.xfer, self.endpoint), session=self._createSession())
			if homedir: return True
		except Exception, e:
			log.warn('%s:%s:%s' % (self.name, e.__class__.__name__, e.message))
			return False

	def getState(self):
		''' Performs availability checks and returns resource state '''
		online     = self.isAlive()
		accessible = self.isAccessible()
		if not online:
			return RES_STATE['offline']
		elif not accessible:
			return RES_STATE['inaccessible']
		else:
			return RES_STATE['online']

	def getAttrs(self):
		''' Returns attribute map for job description '''
		attrs = {}
		attrs['jobbackend']  = self.jobbackend
		attrs['environment'] = self.environment
		attrs['scratchdir']  = self.scratchdir
		return attrs

	def getCacheDir(self):
		if CACHING_ENABLED: return self.filecache.cachedir
	
	def getFromCache(self, filenames):
		if CACHING_ENABLED: return self.filecache.get(filenames)

	def getStats(self):
		return self.stats

	def getXferBackend(self):
		return self.xfer

	def getSagaAuthData(self):
		''' Returns dict for creating SAGA authentication context '''
		method = self.auth['method']

		if method == 'ssh-krb':
			if not self._checkKrbTicket():
				raise RuntimeError('You must have valid kerberos '\
					'ticket for resource %s' % self.name)
		method = 'ssh'
		ctx = {}
		ctx['user_id'] = self.auth['user_id']
		if 'key' in self.auth:
			ctx['user_key']  = '%s.pub' % self.auth['key']
			ctx['user_cert'] = self.auth['key']
		if 'pwd' in self.auth: ctx['user_pass'] = self.auth['pwd']
		return [method, ctx]
	
	def getServiceUrl(self):
		''' Returns URL used for creating saga.Service '''
		proto = None
		method = 'ssh' if self.auth['method'] == 'ssh-krb' else self.auth['method']
		if self.jobbackend != method:
			proto = '%s+%s' % (self.jobbackend, method) if method else self.jobbackend
		else:
			proto = self.jobbackend
		return "%s://%s" % (proto, self.endpoint)

	def getServiceEndpoint(self):
		''' Returns hostname and port of resource '''
		return self.endpoint

	def setStop(self):
		self.stop.set()

	def meetsRequirement(self, key, value):
		''' Checks if resource features contains requirement key and value '''
		if key not in self.features:
			return False

		result = False
		feat   = self.features[key]
		ftype  = type(feat)

		if type(value) == str:
			if ftype == str:
				result = True if feat == value else False
			else:
				result = True if value in feat else False
		elif type(value) == list:
			if ftype == str: return False	
			if len(feat) < len(value): return False

			for item in value:
				if item not in feat: return False
			result = True
		elif type(value) == dict:
			if ftype == str or ftype == list: return False
			for k,v in value.items():
				if k not in feat:
					return False
				result = True if feat[k] == v else False
		return result
	
	def terminate(self):
		super(ResourceAgent, self).terminate()
		self._saveStats()
