#! /usr/bin/env python

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

import time
import hashlib
import sys
import subprocess
import socket, SocketServer
from socket import gethostname
import os
import threading
import logging.config

from tashi.parallel import synchronizedmethod
from tashi.services.ttypes import *
from tashi.util import getConfig, createClient, instantiateImplementation, boolean
from tashi.agents.mauipacket import MauiPacket
import tashi.util

def jobnameToId(jobname):
	return int(jobname.split('.')[-1])

class InstanceHooks():
	def __init__(self, config):
		self.log = logging.getLogger(__file__)
		self.hooks=[]
		items = config.items("MauiWiki")
		items.sort()
		for item in items:
			(name, value) = item
			name = name.lower()
			if (name.startswith("hook")):
				try:
					self.hooks.append(instantiateImplementation(value, config, client, transport, False))
				except:
					self.log.exception("Failed to load hook %s" % (value))
		(self.client, self.transport) = createClient(config)
	def preCreate(self, inst):
		for hook in self.hooks:
			hook.preCreate(inst)
	def postDestroy(self, inst):
		for hook in self.hooks:
			hook.postDestroy(inst)
	def idToInst(self, id):
		instances = self.client.getInstances()
		print 'instances ', instances
		insts = [i for i in instances if str(i.id)==str(id)]
		if len(insts) == 0:
			raise "No instance with ID %s"%id
		if len(insts) > 1:
			raise "Multiple instances with ID %s"%id
		inst = insts[0]
		return inst
	def destroyById(self, id):
		inst = self.idToInst(id)
		self.client.destroyVm(int(id))
		self.postDestroy(inst)
	def activateById(self, id, host):
		inst = self.idToInst(id)
		self.preCreate(inst)
		self.client.activateVm(int(id), host)

def cmplists(a, b):
	for i in range(len(a)):
		if a[i] < b[i]:
			return -1
		if a[i] > b[i]:
			return 1
	return 0

class TashiConnection(threading.Thread):
	def __init__(self, config, client, transport):
		(self.client, self.transport) = createClient(config)

		self.hosts={}
		self.instances={}
		self.users={}

		self.config = config
		self.ihooks = InstanceHooks(config)
		self.log = logging.getLogger(__file__)
		self.refreshTime = float(self.config.get('MauiWiki', 'refreshTime'))
		self.defaultJobTime = str(self.config.get('MauiWiki', 'defaultJobTime'))
		threading.Thread.__init__(self)
		self.daemon = True
	def run(self):
		while True:
			print 'TashiConnection:run updating hosts ...'
			self.updateHosts()
			print 'TashiConnection:run updating instances ...'
			self.updateInstances()
			print 'TashiConnection:run updating users ...'
			self.updateUsers()
			time.sleep(self.refreshTime)
	def wikiHostState(self, host):
		'''Returns a string representing the host state in a form compatable
		with the maui-wiki protocol.  This code simply chooses between
		"Down" and "Unknown":

		- Busy: Node is running some jobs and will not accept additional jobs
		- Down: Resource Manager problems have been detected.  Node is
		incapable of running jobs.
		- Draining: Node is responding but will not accept new jobs
		- Idle: Node is ready to run jobs but currently is not running any.
		Running: Node is running some jobs and will accept additional jobs
		- Unknown: Node is capable of running jobs but the scheduler
		will need to determine if the node state is actually Idle,
		Running, or Busy.'''
		if host.up == False or host.state == HostState.VersionMismatch:
			return "Down"
		if host.state == HostState.Drained:
			return "Draining"
		return "Unknown"

	def wikiInstanceState(self, instance):
		'''Returns a string representing the instance stat in a form compatable
		with the maui-wiki protocol.

		Completed: Job has completed
		Hold: Job is in the queue but is not allowed to run
		Idle: Job is ready to run
		Removed: Job has been canceled or otherwise terminated externally
		Running: Job is currently executing
		Suspended: job has started but execution has temporarily been suspended'''
		tashiToWiki = {InstanceState.Pending:'Idle',
		               InstanceState.Held:'Hold',
		               InstanceState.Exited:'Removed'}
		if tashiToWiki.has_key(instance.state):
			return tashiToWiki[instance.state]
		else:
			return 'Running'

	# Host handling
	def compareHosts(self, host1, host2):
		def ii(host):
			try:
				state = tashi.util.hostStates[host.state]
			except:
				state = 'Unknown'
			return [host.id, host.name, host.up, state, host.memory, host.cores]
		return cmplists(ii(host1), ii(host2))
#	 @synchronizedmethod
	def updateHost(self, host):
		self.hosts[host.id] = host
		self.hosts[host.id].updateTime = time.time()
#	 @synchronizedmethod
	def addHost(self, host):
		self.hosts[host.id] = host
		self.hosts[host.id].updateTime = time.time()
#	 @synchronizedmethod
	def removeHost(self, host):
		self.hosts.pop(host.id)
#	 @synchronizedmethod
	def updateHosts(self):
		if (not self.transport.isOpen):
			self.transport.open()
		hosts = self.client.getHosts()
		for host in hosts:
			if not self.hosts.has_key(host.id):
				self.addHost(host)
			elif self.compareHosts(self.hosts[host.id], host) != 0:
				self.updateHost(host)
		hhosts = {}
		for host in hosts:
			hhosts[host.id] = host
		for host in self.hosts.values():
			if not hhosts.has_key(host.id):
				self.removeHost(host)
	# Instance handling
	def compareInstances(self, instance1, instance2):
		def ii(inst):
			return [inst.id,
			        inst.vmId,
			        inst.hostId,
			        tashi.util.vmStates[inst.state],
			        inst.userId,
			        inst.name,
			        inst.cores,
			        inst.memory,
			        len(inst.disks), # FIXME: this isn't a good way to compare
			        len(inst.nics),   # FIXME: this isn't a good way to compare
			        len(inst.hints)]
		return cmplists(ii(instance1), ii(instance2))
		return 0
	@synchronizedmethod
	def updateInstance(self, instance):
		qt = self.instances[instance.id].queueTime
		self.instances[instance.id] = instance
		self.instances[instance.id].updateTime = time.time()
		self.instances[instance.id].queueTime = qt
	@synchronizedmethod
	def addInstance(self, instance):
		self.instances[instance.id] = instance
		self.instances[instance.id].updateTime = time.time()
		self.instances[instance.id].queueTime = time.time()
	@synchronizedmethod
	def removeInstance(self, instance):
		self.instances[instance.id].state = InstanceState.Exited
		self.ihooks.postDestroy(instance)
	@synchronizedmethod
	def updateInstances(self):
		if (not self.transport.isOpen):
			self.transport.open()
		instances = self.client.getInstances()
		for instance in instances:
			print 'found instance', instance.id
			if not self.instances.has_key(instance.id):
				print "it's a new instance"
				self.addInstance(instance)
			elif self.compareInstances(self.instances[instance.id], instance) != 0:
				self.updateInstance(instance)
		iinsts = {}
		for instance in instances:
			iinsts[instance.id] = instance
		for instance in self.instances.values():
			if instance.state == InstanceState.Exited:
				continue
			if not iinsts.has_key(instance.id):
				print 'removing instance ', instance.id
				self.removeInstance(instance)
	# User handling
	def compareUsers(self, user1, user2):
		if user1.id < user2.id:
			return -1
		elif user1.id > user2.id:
			return 1
		if user1.name < user2.name:
			return -1
		if user1.name > user2.name:
			return 1
		return 0
	@synchronizedmethod
	def updateUser(self, user):
		self.users[user.id] = user
		self.users[user.id].updatetime = time.time()
	@synchronizedmethod
	def addUser(self, user):
		self.users[user.id] = user
		self.users[user.id].updatetime = time.time()
	@synchronizedmethod
	def removeUser(self, user):
		self.users.pop(user.id)
	@synchronizedmethod
	def updateUsers(self):
		if (not self.transport.isOpen):
			self.transport.open()
		users = self.client.getUsers()
		for user in users:
			if not self.users.has_key(user.id):
				self.addUser(user)
			elif self.compareUsers(self.users[user.id], user) != 0:
				self.updateUser(user)
		uusers = {}
		for user in users:
			uusers[user.id] = user
		for user in self.users.values():
			if not uusers.has_key(user.id):
				self.removeUser(user)
	# Get data structures for maui
	# Format is {id:{field:value}}
	@synchronizedmethod
	def getNodes(self, updatetime=0, nodelist=['ALL']):
		if len(nodelist) == 0:
			return {}
		if nodelist[0]=='ALL':
			nodes = [n for n in self.hosts.values() if n.updateTime >= updatetime]
		else:
			nodes = [n for n in self.hosts.values()
					 if n.updateTime >= updatetime and n.name in nodelist]
		nl = {}
		for node in nodes:
			nl[node.name] = {'STATE':self.wikiHostState(node),
			                 'UPDATETIME':str(int(node.updateTime)),
			                 'CPROC':str(node.cores),
			                 'CMEMORY':str(node.memory)}
		return nl
	@synchronizedmethod
	def getJobs(self, updatetime=0, joblist=['ALL']):
		if len(joblist) == 0:
			return {}
		if joblist[0] == 'ALL':
			jobs = [j for j in self.instances.values() if j.updateTime >= updatetime]
		else:
			jobs = [j for j in self.instances.values()
					if j.updateTime >= updatetime and j.id in joblist]
		jl = {}
		for job in jobs:
			id = "%s.%i"%(job.name, job.id)
			jl[id] = {'STATE':self.wikiInstanceState(job),
			          'UNAME':self.users[job.userId].name,
			          'GNAME':self.users[job.userId].name,
			          'UPDATETIME':int(job.updateTime),
			          'QUEUETIME':job.queueTime,
			          'TASKS':'1',
			          'DPROCS':str(job.cores),
			          'DMEM':str(job.memory),
			          'RMEM':str(job.memory),
			          'WCLIMIT':str(self.defaultJobTime)}
			if job.hostId != None:
				jl[id]['TASKLIST'] = self.hosts[job.hostId].name
		return jl
	@synchronizedmethod
	def activateById(self, id, host):
		if not self.instances.has_key(id):
			raise "no such instance"
		self.ihooks.activateById(id, host)
		self.instances[id].state=InstanceState.Activating

class MauiListener(SocketServer.StreamRequestHandler):
	def setup(self):
		global config
		self.log = logging.getLogger(__file__)
		SocketServer.StreamRequestHandler.setup(self)
		self.ihooks = InstanceHooks(config)
		(self.client, self.transport) = createClient(config)
		self.tashiconnection=tashiconnection
		self.auth = config.get('MauiWiki', 'authuser')
		self.key = config.get('MauiWiki', 'authkey')

	def handle(self):
		p = MauiPacket(key=self.key)
		self.istream = self.ostream = self.rfile
		p.readPacket(self.istream)
		self.processPacket(p)

	def processGetNodes(self, p):
		arg = p.data[1]
		arg = arg.split('=')
		arg = arg[1].split(':')
		updatetime = int(arg[0])
		nodelist = arg[1:]
		print 'got GETNODES packet "%s" "%s"'%(updatetime, nodelist)
		r = MauiPacket()
		nodes = tashiconnection.getNodes(updatetime, nodelist)
		numNodes = len(nodes)
		dat = 'ARG=%i#'%numNodes
		first = True
		for node, attributes in nodes.iteritems():
			if first:
				dat = dat + '%s:'%node
				first = False
			else:
				dat = dat + '#%s:'%node
			attrs = ['%s=%s'%(a,v) for a,v in attributes.iteritems()]
			dat = dat + ';'.join(attrs)+';'
		r.set(['SC=0', dat], auth=self.auth, key=self.key)
		return r

	def processGetJobs(self, p):
		arg = p.data[1]
		arg = arg.split('=')
		arg = arg[1].split(':')
		updatetime = int(arg[0])
		joblist = arg[1:]
		print 'got GETJOBS packet "%s" "%s"'%(updatetime, joblist)
		r = MauiPacket();
		jobs = tashiconnection.getJobs(updatetime, joblist)
		numJobs = len(jobs)
		dat = 'ARG=%i#'%numJobs
		first = True
		for job, attributes in jobs.iteritems():
			if first:
				dat = dat +'%s:'%job
				first = False
			else:
				dat = dat +'#%s:'%job
			# FIXME: support limits
			attributes['WCLIMIT']=str(10000)
			attrs = ['%s=%s'%(a,v) for a,v in attributes.iteritems()]
			dat = dat + ';'.join(attrs) + ';'
		r.set(['SC=0', dat], auth=self.auth, key=self.key)
		return r

	def processStartJob(self, p):
		job = p.data[1]
		job = job.split('=')[1].strip()
		job = job.split('.')[-1]
		tasklist = p.data[2].split('=')[1].split(':')
		print 'STARTJOB ', job, tasklist
		try:
			hosts = self.client.getHosts()
			print 'hosts ', hosts
			host = [h for h in hosts if h.name == tasklist[0]][0]
			self.tashiconnection.activateById(jobnameToId(job), host)
			print '\tactivated VM!'
			r = MauiPacket()
			r.set(['SC=0','RESPONSE=VM %s started on host %s'%(job, tasklist[0])])
			return r
		except Exception, e:
			# FIXME: make this a real failure response
			print 'Oh noes! ', e
			r = MauiPacket()
			r.set(['SC=-1', 'RESPONSE=%s'%str(e)])
			return r

	def processCancelJob(self, p):
		job = p.data[1]
		job = job.split('=')[1].strip()
		print 'CANCELJOB ', job
		try:
			self.client.destroyVm(jobnameToId(job))
			print '\tdestroyed VM!'
			r = MauiPacket()
			r.set(['SC=0', 'RESPONSE=VM %s destroyed'%job])
			return r
		except Exception, e:
			# FIXME: make this a real failure response
			print 'Oh noes! ', e
			r = MauiPacket()
			r.set(['SC=-1', 'RESPONSE=%s'%str(e)])
			return r

	def processPacket(self,p):
		dat = p.data
		if not p.verifyChecksum():
			print p
			print 'bad checksum'
			return
		r = None
		if dat[0] == 'CMD=GETNODES':
			r = self.processGetNodes(p)
		elif dat[0] == 'CMD=GETJOBS':
			r = self.processGetJobs(p)
		elif dat[0] == 'CMD=STARTJOB':
			r = self.processStartJob(p)
		elif dat[0] == 'CMD=CANCELJOB':
			r = self.processCancelJob(p)
		else:
			print 'got unknown packet'
			print p.prettyString()
			r = MauiPacket()
			r.set(['SC=-810', 'RESPONSE=command not supported'])
		print r.prettyString()
		self.ostream.write(str(r))
		self.ostream.flush()
		self.ostream.close()
		self.ostream.close()
		self.istream.close()

if __name__ == '__main__':
	(config, configFiles) = getConfig(["Agent"])
	publisher = instantiateImplementation(config.get("Agent", "publisher"), config)
	tashi.publisher = publisher
	(client, transport) = createClient(config)
	logging.config.fileConfig(configFiles)
	tashiconnection = TashiConnection(config, client, transport)
	tashiconnection.start()

	HOST, PORT = '', 1717
	server = SocketServer.TCPServer((HOST,PORT), MauiListener)
	server.serve_forever()