import os,sys
import re
import copy
from py_autoflow.queue_manager import QueueManager

class Slurm2Manager(QueueManager):
	# SLURM 20 or greater
	priority = 100

	@classmethod
	def isavailable(cls, options):
		available = False
		shell_output = cls.system_call_class("type 'srun'")
		if options.get('show_submit_command'): print(shell_output, file=sys.stdout)
		if len(shell_output) > 0 :
			shell_output = cls.system_call_class("srun --version", show_cmd = options.get('show_submit_command'))
			if options.get('show_submit_command'): print(shell_output, file=sys.stdout)
			version_string = re.search(r'slurm ([0-9\.]+)', shell_output).groups()[0] # "slurm 17.11.4"
			slurm_version = int(version_string.split('.')[0])
			if slurm_version >= 20: available = True 
		return available


	def write_header(self, id, job, sh_name):
		if not job.attrib['ntask']:
			self.write_file(sh_name, f"#SBATCH --cpus-per-task={job.attrib['cpu']}")
		else:
			self.write_file(sh_name, f"#SBATCH --ntasks={job.attrib['cpu']}")
			if job.attrib['multinode'] > 0: self.write_file(sh_name, f"#SBATCH --nodes={job.attrib['multinode']}") 
		self.write_file(sh_name, f"#SBATCH --mem={job.attrib['mem']}")
		self.write_file(sh_name, f"#SBATCH --time={job.attrib['time']}")
		if job.attrib.get('node') != None: self.write_file(sh_name, f"#SBATCH --constraint={job.attrib['node']}") 
		self.write_file(sh_name, f'#SBATCH --error=job.%J.err')
		self.write_file(sh_name, f'#SBATCH --output=job.%J.out')
		for param, val in job.each_add_opt():
			if val != None:  # param value queue option
				self.write_file(sh_name, f"#SBATCH --{param}={self.parse_additional_options(val, job.attrib)}")
			else:  # boolean queue option
				self.write_file(sh_name, f"#SBATCH --{param}")
		if job.attrib['ntask'] and job.attrib['cpu_asign'] == 'list': self.write_file(sh_name, 'srun hostname -s > workers') 

	def parse_additional_options(self, string, attribs):
		expresions = ['%C', '%T', '%M' '%N' ]
		values = [attribs['cpu'], attribs['time'], attribs['mem'], attribs['node']]
		new_string = copy.copy(string)
		for i,exp in enumerate(expresions): new_string = re.sub(exp, str(values[i]), new_string)
		return new_string

	def submit_job(self, job, ar_dependencies):
		final_dep = self.get_all_deps(ar_dependencies)
		dependencies = None
		if len(final_dep) > 0: 
			dependencies ='--dependency=afterok:' + (':').join(final_dep)
		else:
			dependencies = ''
		cmd = f"sbatch {dependencies} {job.name}.sh"
		if self.show_submit: print(cmd, file=sys.stdout)
		if not self.verbose: #Append submission data to job
			queue_id = self.get_queue_system_id(self.system_call(cmd, job.attrib['exec_folder']))
		else:
			queue_id = None
		return queue_id

	def get_queue_system_id(self, shell_output):
		queue_id = None
		submission_match = re.search(r'Submitted batch job (\d+)', shell_output)
		if submission_match == None: raise Exception(f"A queue id cannot be obtained. The queue manager has given this message:{shell_output}") 
		queue_id = submission_match.groups()[0]
		return queue_id

	def get_queue_system_dependencies(self, ar_dependencies):
		queue_system_ids = [ self.commands[dependency].queue_id for dependency in ar_dependencies ]
		return queue_system_ids

	def get_all_deps(self, ar_dependencies):
		final_dep = []
		if len(ar_dependencies) > 0: final_dep.extend(self.get_queue_system_dependencies(ar_dependencies)) 
		final_dep.extend(self.external_dependencies)
		final_dep = [ x for x in final_dep if x != None ]
		return final_dep
