import sys, os, re
import json
from termcolor import colored

from py_autoflow.batch import Batch
from py_autoflow.wflow_smith import WfSmith

class Stack:
	def __init__(self, exec_folder, options):
		self.set_initialize(exec_folder, options = options)

	def set_initialize(self, exec_folder, options= None): # method to initialize and to clean the object
		self.commands = {}
		self.variables = {}
		self.exec_folder = exec_folder #TODO move this to queue_manager
		self.resources = {}
		self.wf_provider = None
		if options == None:
			self.do_retry = False
			self.options = {}
			self.workflow = ''
			self.external_variables = []
			self.white_list = []
			self.black_list = []
		else:
			self.do_retry = options.get('retry')
			self.options = options
			self.workflow = options.get('workflow')
			self.external_variables= options.get('Variables')
			self.white_list = options.get('white_list')
			self.black_list = options.get('black_list')
		self.jobs = {}

	def clean(self):
		self.set_initialize(None)

	def inspect(self):
		for id, job in self.jobs.items(): print(job.inspect(all_attribs = False))

	def parse_resources(self, options):
		default_opts = {
			'cpu' : options.get('cpus'), 
			'mem' : options.get('memory'), 
			'time' : options.get('time'),
			'node' : options.get('node_type'),
			'multinode' : options.get('use_multinode'),
			'ntask' : options.get('use_ntasks'),
			'additional_job_options' : options.get('additional_job_options')
		}
		wf_manager = None
		resources = {}
		resources_file = options.get('resources')
		if resources_file != None:
			resources_data = None
			with open(resources_file, 'r') as f: resources_data = json.load(f)
			resources_profiles = resources_data.get('resources')
			#{ 'default' : [
			#       'cpu': int,
			#       'node_type': string,
			#       'memory': string
			#       'time': string
			#		'ntask': bool => use more than one node or not
			#       'multinode': int => 0 to no limit in the number of nodos or N to require all cpus in N nodes
			#]}
			if resources_profiles != None and resources_profiles.get('default') != None: 
				default_opts.update(resources_profiles['default']) # Merge defined options in file for flag options, maybe the file do NOT define all the necessary options for default profile
			resources_profiles['default'] = default_opts
			resources = resources_profiles
			virtualization_profiles = resources_data.get('virt')
			if virtualization_profiles != None:
				paths = [os.getenv("AUTOFLOW_GENERAL", None), os.getenv("AUTOFLOW_USER", None)]
				paths = [p for p in paths if p != None]
				if len(paths) == 0 or options.get('res_in_wf'): paths = [options['output']]
				wf_manager = WfSmith(paths)
				wf_manager.load_available_virtualizations()
				wf_manager.configure_virt_from_resource_profiles(virtualization_profiles, paths[-1])
				self.wf_manager = wf_manager
		else:
			resources['default'] = default_opts

		Batch.general_computation_attrib = resources

	def parse(self):
		self.clean_template()

		variables_lines = []
		node_lines = []

		node_beg = False		
		for line in self.workflow.splitlines():
			if '{' in line: node_beg = True	# This check the context of a variable
			if '}' in line:					# if a variable is within a node,
				if node_beg:				# we consider tha is a bash variable not a static autoflow variable
					node_beg = False
				else:
					node_beg = True
			if not node_beg and re.search(r"^\$", line):
				variables_lines.append(line)
			else:
				node_lines.append(line)
		self.load_variables(variables_lines, self.variables)
		self.load_variables(self.external_variables, self.variables)
		self.parse_nodes(node_lines)

	def clean_template(self):
		wf = self.workflow.replace(r"\#.+$",'')	#Delete comments
		wf = wf.replace("\t",'')		#Drop tabs
		wf = wf.replace(r"\n+","\n")	#Drop empty lines
		self.workflow = wf.replace(r"^\s*",'')

	def load_variables(self, variables_lines, variable_type):
		if variables_lines != None:
			for line in variables_lines:
				if isinstance(line, str):
					pairs = line.rstrip().replace(r"\s",'').split(',')
					for pair in pairs:
						fields = pair.split('=', 1)
						if len(fields) == 2:
							var, value = fields
							variable_type[var] = value
						else:
							self.load_var_file(fields[0], variable_type)
				else:
					if len(line) == 2:
						variable_type[line[0]] = line[1]
					else:
						self.load_var_file(line[0], variable_type)

	def load_var_file(self, path, variable_type):
		if os.path.exists(path):
			with open(path) as f:
				for line in f:
					if re.match(r"^#", line) or len(line) == 0: continue
					var, value = line.rstrip().replace(r"\s",'').split('=')
					variable_type[var] = value
		else:
			raise Exception(f"Variable file {path} not exists")

	def scan_nodes(self, execution_lines):
		template_executions = '\n'.join(execution_lines)
		template_executions = self.replace_variables(template_executions)
		executions = [] # tag, initialize, main_command
		states = {} #name => [state, id(position)]
					#t => tag, i => initialize , c => command
		open_nodes = []
		for line in template_executions.splitlines():
			line = line.strip() #Clean al whitespaces at beginning and the end of string
			if line == '': continue
			if len(open_nodes) > 0: node = states[open_nodes[-1]]
			
			# Create nodes and asign nodes states
			#----------------------------------------
			node_match = re.search(r"(\S*\)){$", line)
			if node_match != None: #Check tag and create node
				name = node_match.groups()[0]
				executions.append([name, '', '']) # create node
				states[name] = ['i', len(executions) - 1]
				open_nodes.append(name)
			elif line == '?': #Check command
				node[0] = 'd'
			elif line == '}': #Close node
				finished_node = open_nodes.pop()
				if len(open_nodes) > 0:
					parent_node = states[open_nodes[-1]][-1] #position
					child_node = states[finished_node][-1]
					parent_execution = executions[parent_node]
					if isinstance(parent_execution[2], str):
						parent_execution[2] = [child_node]
					else:
						parent_execution[2].append(child_node)
			# Add lines to nodes
			#----------------------
			elif states[open_nodes[-1]][0] == 'i': #Add initialize line
				executions[node[-1]][1] = executions[node[-1]][1] + line +"\n"
			elif states[open_nodes[-1]][0] == 'd': #Add command line
				executions[node[-1]][2] = executions[node[-1]][2] + line +"\n"

		return executions

	def replace_variables(self, string):
		variables = [ [name, value] for name, value in self.variables.items() ]
		variables.sort(key=lambda x: len(x[0]), reverse=True)
		for name, value in variables:
			string = string.replace(name, value)
		return string

	def parse_nodes(self, execution_lines):
		dinamic_variables = []
		nodes = self.scan_nodes(execution_lines)
		self.create_ids(nodes)
		Batch.init_indexes() # To clean start job relations indexes
		for tag, init, command, index in nodes: #Takes the info of each node of workflow to create the job
			new_batch = Batch(tag, init, command, index, self.exec_folder)
			#dinamic_variables = dinamic_variables + new_batch.handle_dependencies(dinamic_variables)
			self.commands[new_batch.name] = new_batch

		# link each parent batch to a child batch
		for name, batch in self.commands.items(): batch.asign_child_batch()

	def create_ids(self, nodes):
		for i, node in enumerate(nodes): node.append(i)

	def get_jobs_relations(self):
		relations = {}
		for name, job in self.get_jobs(): 
			relations[name] = job
			self.check_authorized_job(job)
		self.jobs = relations
		return relations

	def check_authorized_job(self, job):
		done = None
		if len(self.black_list) > 0:
			if any([ re.match(regexp, job.name) for regexp in self.black_list]) : 
				done = True
			else:
				done = False
		if len(self.white_list) > 0:
			if any([ re.match(regexp, job.name) for regexp in self.white_list]) : 
				done = False
			else:
				done = True
		if done != None: job.attrib['done'] = done

	def get_jobs(self):
		jobs =[]
		for name, batch  in self.commands.items():
			if batch.has_jobs(): continue #parent batch (intermediates)
			for j in batch.get_jobs():
				#folder = self.asign_folder(j) #TODO move this to queue_manager
				jobs.append([j.name, j])

		# This loop needs all defined jobs to set the dependencies path but if removed, could be added to previous loop
		for job in jobs: #TODO move this to queue_manager
			j_name, j = job
			#self.set_dependencies_path(job)
			job[0] = j_name.replace(')','') #Clean function characters on name
			job[1].name = j.name.replace(')','') #Clean function characters on name
		return jobs

