# simlib.py -- shared necessary library functions across all sim-* utilties
#
#
#
#

import os, sys, re
import libutil
import socket
import traceback

class SimLib:
	def __init__(self, env):
		# we need SimEnvironemtn
		self.SimEnvironment = env
		self.OptionsManager = self.SimEnvironment.OptionsManager
		self.ConfigurationDatabase = self.SimEnvironment.ConfigurationDatabase
		self.DefineDatabase = self.SimEnvironment.DefineDatabase
	
	def GetHostNameAlias(self):
		
		# first check if our hostname was specified via an option
		if self.OptionsManager.HasOption('hostname'):
			return self.OptionsManager.GetOption('hostname')
		
		# get the users home directory, check for ~/.hostname
		homedir = os.environ['HOME']
		home_hostname = "%s%s%s" % (homedir, os.sep, ".hostname")
		if libutil.FileExists(home_hostname):
			hostname = self.GetFileContents(home_hostname).strip()
			return hostname
		
		# lets get the hostname
		hostname = os.popen("hostname").read().strip()
		
		try:
			(name, aliases, _) = socket.gethostbyname_ex(hostname)
		except:
			print "Warning: could not resolve hostname %s" % hostname
			return hostname
		
		
		
		# out of all our names, find one that has a dot, if possible.
		names = aliases
		names.append(hostname)
		names.append(name)
		
		for n in names:
			if n.count("."):
				return n

		if len(name) > 0:
			return name
		
		return hostname
	
	def BoolToYesNo(self, val):
	
		if val == True:
			return 'yes'
		
		return 'no'
		
	
	def GetLocalEnvironment(self):
		
		machine = self.SimEnvironment.LocalMachine
		machineEntry = self.SimEnvironment.LocalMachineEntry		
		
		req_keys = ['user', 'sourcebasedir']
		
		self.VerifyKeys(machineEntry, req_keys)
		
		sourceBaseDir = self.GetLocalSourceBaseDir()
		path = self.GetDirSuffix(sourceBaseDir)
		
		return (machine, machineEntry, sourceBaseDir, path)
		
	def GetMachineName(self, machineName=None):
		if machineName == None:
			machineName = self.GetHostNameAlias()
		
		machines = self.ConfigurationDatabase.GetMachines()
		
		found_machines = []
		
		for m in machines:
			if m == machineName:
				found_machines.append(m)
			else:
				entry = self.ConfigurationDatabase.GetMachine(m)

				if entry.HasKey('hostname') and entry.hostname == machineName:
					found_machines.append(m)
				else:
					if entry.HasKey("aliaspattern"):
						if libutil.Matches(entry.aliaspattern, machineName):
							found_machines.append(m)
		
		error = None
		if len(found_machines) == 0:
			error = "Unknown machine name %s" % machineName

		if len(found_machines) > 1:
			error = "Could not identify machine %s -- possible matches are %s" % (machineName, found_machines)
		
		if error:
			print error
			sys.exit(1)
		
		return found_machines[0]
	
	# Get source base directory
	def GetLocalSourceBaseDir(self):
	
		if self.OptionsManager.HasOption('sourcebasedir'):
			sourcebasedir = self.OptionsManager.GetOption('sourcebasedir')
		else:
			return self.GetSourceBaseDir(self.SimEnvironment.LocalMachineEntry)
			
		if not(sourcebasedir[0] == "/"):
			cwd = os.getcwd()
			sourcebasedir = "%s%s%s" % (cwd, os.sep, sourcebasedir)
		
		if not(os.path.exists(sourcebasedir)):
			print 'Cannot access source base directory "%s"' % sourcebasedir
		
		return sourcebasedir
	
	def BaseName(self, filename):
		(path, file) = os.path.split(filename)
		return file
	
	def FileBaseName(self, filename):
		(file, ext) = self.BaseName(filename).split(".")
		return file
		
	def GetBaseDir(self, machineEntry):
		
		if self.OptionsManager.HasOption('basedir'):
			basedir = self.OptionsManager.GetOption('basedir')
		else:
			machineEntry = self.SimEnvironment.LocalMachineEntry
			
			self.VerifyKeys(machineEntry, ['user', 'basedir'])
			
			user = machineEntry.user
		
			self.DefineDatabase.Set('USER', user)
			basedir = self.DefineDatabase.SubAll(machineEntry.basedir)
			
		if not(basedir[0] == "/"):
			cwd = os.getcwd()
			basedir = "%s%s%s" % (cwd, os.sep, basedir)
		
		if not(os.path.exists(basedir)):
			print 'Cannot access source base directory "%s"' % basedir
		
		return basedir
	
	def GetSourceBaseDir(self, machineEntry):

		if not(machineEntry.HasKey('user')):
			print "Error, machine entry %s does not have a user defined" % machineName
			sys.exit(1)
		
		user = machineEntry.user
		
		self.DefineDatabase.Set('USER', user)
		sourcebasedir = self.DefineDatabase.SubAll(machineEntry.sourcebasedir)
			
		if not(sourcebasedir[0] == "/"):
			cwd = os.getcwd()
			sourcebasedir = "%s%s%s" % (cwd, os.sep, sourcebasedir)
		
		if not(os.path.exists(sourcebasedir)):
			print 'Cannot access source base directory "%s"' % sourcebasedir
		
		return sourcebasedir
	
	def GetDirSuffix(self, prefix):
		
		dir = os.getcwd()
		os.chdir(prefix)
		
		real_prefix = os.getcwd()
		os.chdir(dir)
		
		pattern = "^%s/(.*)$" % real_prefix
		
		p = re.compile(pattern)
		matches = p.match(dir)
		
		if matches == None:
			error = "Called from the wrong location."
			error = "%s\n%s" % (error, "Current directory is '%s'" % dir)
			error = "%s\n%s" % (error, "but expected a subdirectory of '%s'." % prefix)
			error = "%s\n%s" % (error, "It is also be possible that you need to correct your 'sourcebasedir' entry")
			error = "%s\n%s" % (error, "in the mdb entry for this machine.")
			print error
			sys.exit(1)
		
		return matches.group(1)
	
	def GetCactusPath(self):
	
		basepath = self.SimEnvironment.BASE_PATH
		
		if basepath.count("/py"):
			pops = 2
		else:
			pops = 1
		
		parts = basepath.split(os.sep)
		
		for i in range(pops):
			parts.pop()
		
		path = os.sep.join(parts)
		return path
				
	
	def GetUsername(self):
		
		if "USER" not in os.environ:
			username = os.popen("whoami").read()
		else:
			username = os.environ["USER"]
		
		return username
		
	def GetSSHCommand(self, machine, cmd, options=None):
		
		if options == None:
			options = ''
		
		currentMachine = machine
		
		while currentMachine != None:
			entry = self.ConfigurationDatabase.GetMachine(currentMachine)
			
			req_key = ['sshcmd', 'sshopts', 'user', 'hostname']
			self.VerifyKeys(entry, req_key)
			
			sshcmd = entry.sshcmd
			sshopts = entry.sshopts
			user = entry.user
			hostname = entry.hostname
			
			self.DefineDatabase.Set('USER', user)
			
			ssh = "%s %s %s %s@%s" % (sshcmd, sshopts, options, user, hostname)
			ssh = self.DefineDatabase.SubAll(ssh)
			
			if entry.HasKey("sshsetup"):
				cmd = "{ %s; } && %s" % (entry.sshsetup, cmd)
			
			cmd = self.QuoteSafe(cmd)
			cmd = "%s %s" % (ssh, cmd)
			
			trampoline = self.GetTrampoline(currentMachine)
			
			if trampoline != None:
				if not(self.ConfigurationDatabase.HasMachineEntry(trampoline)):
					print "Unknown trampoline %s specified for machine %s" % (trampoline, currentMachine)
					sys.exit(1)
				
				tentry = self.ConfigurationDatabase.GetMachine(trampoline)
				if tentry.HasKey("localsshsetup"):
					cmd = "{ %s; } && %s" % (self.DefineDatabase.SubAll(tentry.localsshsetup), cmd)
				
			
			currentMachine = trampoline
		
		return cmd
	
	def ExecuteCommand(self, command):
		print "Executing: %s\n" % command

		environcmd = self.GetMachineOption('environcmd')
		command = "{ %s; } && %s" % (environcmd, command)
		
		ret = os.system(command)
		
		if ret > 0:
			print "Error %s occured while executing command \"%s\"" % (ret, command)
			
	def GetArguments(self):
		args = []
		if self.OptionsManager.HasOption("argument"):
			args.append(self.OptionsManager.GetOption('argument'))
		
		return args
		
	def QuoteSafe(self, cmd):
		return "'%s'" % cmd.replace("'", r"'\''")
		
	def VerifyKeys(self, machineEntry, req_keys):
		for key in req_keys:
			if not(machineEntry.HasKey(key)):
				print "Error, machine %s is missing a required key: %s" % (machineEntry.EntryName, key)
				traceback.print_tb(sys.last_traceback)
				sys.exit(1)
		
	def GetIOMachine(self, machineName):
		return self.GetMachineByKey(machineName, 'iomachine')
	
	def GetTrampoline(self, machineName):
		return self.GetMachineByKey(machineName, 'trampoline', True)
		
	def GetMachineByKey(self, machineName, key, useNone=False):
		entry = self.ConfigurationDatabase.GetMachine(machineName)
		
		if not(entry.HasKey(key)):
			if useNone:
				return None
			else:
				return machineName

		mm = entry.GetKey(key)
		
		if not(self.ConfigurationDatabase.HasMachineEntry(mm)):
			print "Error, specified %s %s for machine %s does not exist" % (key, mm, machineName)
			sys.exit(1)
		
		return mm
	
	def GetUseNodes(self):
		
		if not(self.OptionsManager.HasOption('use-nodes')):
			return []
		
		nodes = self.OptionsManager.GetOption('use-nodes')
		
		return nodes.split(",")
		
	def GetDefaultConfiguration(self):
		OptionsManager = self.OptionsManager
		ConfigurationDatabase = self.ConfigurationDatabase
		
		build_debug = OptionsManager.GetOption("debug")
		build_optimise = OptionsManager.GetOption("optimise")
		build_unsafe = OptionsManager.GetOption("unsafe")
		build_profile = OptionsManager.GetOption("profile")
		
		build_optimise = not(build_debug)
		
		dEntry = ConfigurationDatabase.GetConfigOption("default-configuration-name");
		if dEntry == None:
			configName = "sim"
		else:
			configName = dEntry
			
		config = configName
		
		if build_debug:
			config = "%s-%s" % (config, "debug")
		
		if build_optimise:
			rr = "optimise"
		else:
			rr = "nooptimise"
		
		config = "%s-%s" % (config, rr)
		
		if build_unsafe:
			config = "%s-%s" % (config, "unsafe")
	
		if build_profile:
			config = "%s-%s" % (config, "profile")
		
		return config
		
	def FileExists(self, filename):
		return os.path.exists(filename)
		
	def GetFileContents(self, filename, default="", suppressWarnings=True):
		if filename == None:
			return default

		try:
			fptr = open(filename, "r")
			contents = fptr.read()
			fptr.close()
			return contents
		except:
			if suppressWarnings == False:
				print "Warning: could not open %s for reading" % filename
			
			return default
	
	def WriteContents(self, filename, contents):
		try:
			fptr = open(filename, "w")
		except:
			print "Error: could not open %s for writing" % filename
			sys.exit(1)
		
		fptr.write(contents)
		fptr.close()
		
	
	def RemoveFile(self, filename):
		try:
			os.unlink(filename)
		except OSError:
			return
	
	def RenameFile(self, old, new):
		try:
			os.rename(old, new)
		except OSError:
			return
	
	def GetVersion(self, fileContents):
		regex = 'VERSION\s*=\s*([^#;\n]*)[#;]?.*'
		rr = re.compile(regex)
		m = rr.search(fileContents)
		if m != None:
			return m.group(1)
		else:
			return "0000-00-00"
		
		
	def BuildPath(self, args):		
		return os.sep.join(args)
	
	def GetOptionsList(self, required=False):
		return self.GetEtcFile("optionlist", required)
	
	def GetThornList(self, required=False):
		return self.GetEtcFile("thornlist", required)
		
	def GetParFile(self, required=False):
		return self.GetEtcFile("parfile", required)
	
	def GetSubmitScript(self, required=False):
		return self.GetEtcFile("submitscript", required)
		
		
	def GetMachineOption(self, option, defaultValue=None):
		
		value = None
		
		if self.OptionsManager.HasOption(option):
			value = self.OptionsManager.GetOption(option)
			return value
		else:
			value = defaultValue
		
		if value == None:

			# make sure option is set
			self.VerifyKeys(self.SimEnvironment.LocalMachineEntry, [option])
			
			value = self.SimEnvironment.LocalMachineEntry.GetKey(option)
		
		return value
	
	def GetProcs(self, existingProperties, procs_arg=None):
		
		machineEntry = self.SimEnvironment.LocalMachineEntry
		
		minppn = int(machineEntry.GetKey("min-ppn"))
		maxppn = int(machineEntry.GetKey("ppn"))
		maxnodes = int(machineEntry.GetKey("nodes"))
		
		procs = None
		
		if self.OptionsManager.HasOption('procs'):
			procs = self.OptionsManager.GetOption('procs')
		
		if procs == None and procs_arg == None:
			print "Error, number of processors to use in the simulation was not specified"
			sys.exit(1)
		
		if procs_arg != None:
			procs = procs_arg
		
		try:
			procs = int(procs)
		except ValueError:
			print "couldn't coerse procs value %s to an integer" % procs
			sys.exit(1)
			
		if procs < 1:
			print "Illegal number of processors specified for the simulation"
			sys.exit(1)
		
		ppn = int(self.GetMachineOption('ppn'))
		
		if ppn == None:
			ppn = maxppn
			
			if ppn > procs:
				ppn = procs
		
		#if (defined $minppn && ($ppn < $minppn || $ppn > $maxppn)) {
        #die "Illegal number of requested processors per node specified: specified ppn=$ppn (min-ppn is $minppn, max-ppn is $maxppn)";
		
		if minppn != None and (ppn < minppn or ppn > maxppn):
			print "Illegal number of requested processors per node specified: specified ppn=%s (min-ppn is %s, max-ppn is %s)" % (ppn, minppn, maxppn)
			sys.exit(1)
		
		ppnused = None
		
		if self.OptionsManager.HasOption('ppn-used'):
			ppnused = self.OptionsManager.GetOption('ppn-used')
		
		if ppnused == None:
			ppnused = ppn
			
		try:
			ppnused = int(ppnused)
		except ValueError:
			print "couldn't coerse ppn-used value %s to an integer" % ppnused
			sys.exit(1)
		
		if ppnused < 1:
			print "Illegal number of used processors per node specified: specified ppn-used=%s" % ppnused
			sys.exit(1)
		
		if ppnused > ppn:
			print "WARNING: Too many used processors per node specified: specified ppn-used=%s (ppn is %s)" % (ppnused, ppn)
		
		if existingProperties != None:
			num_threads = self.GetMachineOption('num-threads', existingProperties.numthreads)
		else:
			num_threads = self.GetMachineOption('num-threads', 1)
		
		if num_threads < 1:
			print "Illegal number of threads per process specified: specified num-threads=%s" % num_threads
			sys.exit(1)
		
		if num_threads > procs:
			print "WARNING: Too many threads per process specified: specified num-threads=%s (ppn-used is %s)" % (num_threads, ppnused)
		
		num_procs = procs / int(num_threads)
		
		if num_procs < 1:
			num_procs = 1
		
		nodes = int((procs + ppnused - 1) / ppnused)
		
		if nodes < 1:
			nodes = 1
			
		if nodes > maxnodes:
			print "Too many nodes specified: nodes=%s (maxnodes is %s)" % (nodes, maxnodes)
		
		procs_requested = nodes * ppn
		
		return (nodes, ppnused, procs, ppn, procs_requested, num_procs, num_threads)
		
	def GetEtcFile(self, option, required=False):
		if self.OptionsManager.HasOption(option):
			ff = self.OptionsManager.GetOption(option)
			if not(os.path.exists(ff)):
				print "Error: specified %s %s does not exist or is not readable" % (option, ff)
				sys.exit(1)
			
			return os.path.abspath(ff)
		
		if required == False:
			return None

		machineEntry = self.SimEnvironment.LocalMachineEntry
		
		# make sure option is set
		self.VerifyKeys(machineEntry, [option])
		
		ff = machineEntry.GetKey(option)
		
		if ff[0] != os.sep:
			folder = "%ss" % option
			eff = self.BuildPath((self.SimEnvironment.ETC_PATH,folder,ff))

			if os.path.exists(eff):
				ff = eff
		
		if not(os.path.exists(ff)):
			print "Error: specified %s %s does not exist or is not readable" % (option, ff)
			sys.exit(1)
		
		return ff