import os
import glob
import warnings
import sys
import pathlib
from pydoc import importfile
import time
import json
import hashlib
import subprocess
import re

from py_autoflow.logging import *

class QueueManager:
	priority = -1

	########################################################################################
	## SELECT AND PREPARE MANAGER
	########################################################################################
	@classmethod	
	def isavailable(cls, options):
		return False

	@classmethod
	def system_call_class(cls, cmd, path = None, show_cmd = False, remote = False, ssh = None):
		if path != None: cmd = f"cd {path}; " + cmd
		if show_cmd: sys.stdout.print(cmd)
		if remote:
			raise Exception("Detected remote execution: is not implemented")
			#call = @ssh.exec!(cmd)
		else:
			call = subprocess.check_output(cmd, shell = True).decode("utf-8")
		return call

	@classmethod
	def select_queue_manager(cls, exec_folder, options, jobs):
		if options['batch']: 
			from py_autoflow.queue_managers.bash_manager import BashManager
			queue_manager = BashManager
		else:
			path_managers = os.path.join(os.path.dirname(__file__),'queue_managers')
			for manager in glob.glob(path_managers+'/*.py'): importfile(manager)
			queue_manager = cls.select_manager(options)
		warnings.warn(f"Selected queue manager: {queue_manager}")
		return queue_manager(exec_folder, options, jobs)

	@classmethod
	def select_manager(cls, options):
		queue_manager = None
		priority = 0
		for descendant in QueueManager.__subclasses__():
			print(f"{descendant}")
			if descendant.isavailable(options) and priority <= descendant.priority:
				queue_manager = descendant
				priority = descendant.priority
		return queue_manager

	########################################################################################
	## USE MANAGER
	########################################################################################
	def __init__(self, exec_folder, options, commands):
		self.set_initialize(exec_folder, options, commands)

	def set_initialize(self, exec_folder = 'exec', options = None, commands = {}): # method to initialize and to clean the object
		self.folder_name = 'program_name'
		self.exec_folder = exec_folder
		self.raw_folder = os.path.join(exec_folder, '_file')
		self.commands = commands
		self.files = {}
		self.active_jobs = []
		self.wf_smith = None
		if options == None:
			self.verbose = 0
			self.show_submit = False
			self.job_identifier = ''
			self.remote = False
			self.ssh = None
			self.asign_job_folder = True
			self.write_sh = True
			self.external_dependencies = []
			self.extended_logging = False
			self.comment = False
			self.sleep_time = 0
			self.linked_folders = False
		else:
			if options['key_name']: self.folder_name = 'job_name' 
			self.verbose = options.get('verbose')
			self.show_submit = options.get('show_submit_command')
			self.job_identifier = options.get('identifier')
			self.remote = options.get('remote')
			self.ssh = options.get('ssh')
			self.asign_job_folder = options.get('asign_job_folder')
			self.write_sh = options.get('write_sh')
			self.external_dependencies = options.get('external_dependencies')
			self.extended_logging = options.get('extended_logging')
			self.comment = options.get('comment')
			self.sleep_time = options.get('sleep_time')
			self.linked_folders = options.get('linked_folders')
		self.additional_job_options = []

	def clean(self):
		self.set_initialize()


	## FILE AND CMD MANAGING
	########################################################################################

	def create_file(self, file_name, path):
		self.files[file_name] = [path, '']

	def write_file(self, file_name, content):
		file = self.files[file_name]
		file[-1] = file[-1] + content +"\n"

	def close_file(self, file_name, permissions = None): # use octal representation for permissions (four digits, 0-7 with prefix 0o if a 0 is necessary, i.e: 0o755)
		path, content = self.files.pop(file_name)
		file_path = os.path.join(path, file_name)
		if self.remote:
			raise Exception("Detected remote execution: is not implemented")
			#@ssh.exec!("echo '#{content}' > #{file_path}")
			#@ssh.exec!("chmod #{permissions} #{file_path}") if !permissions.nil?
		else:
			with open(file_path, 'w') as local_file: local_file.write(content)
			if permissions != None: os.chmod(file_path, permissions) 

	def create_folder(self, folder_name):
		if self.remote:
			raise Exception("Detected remote execution: is not implemented")
			#@ssh.exec!("if ! [ -d  #{folder_name} ]; then mkdir -p #{folder_name}; fi")
		else:
			if not os.path.exists(folder_name): os.mkdir(folder_name)

	def system_call(self, cmd, path = None, show_cmd = False):
		if path != None: cmd = f"cd {path}; " + cmd
		if show_cmd: sys.stdout.print(cmd)
		if self.remote:
			raise Exception("Detected remote execution: is not implemented")
			#call = @ssh.exec!(cmd)
		else:
			call = subprocess.check_output(cmd, shell = True).decode("utf-8")
		return call

	## EXECUTING WORKFLOW WITH MANAGER
	########################################################################################

	def exec(self, wf_provider = None):
		self.wf_provider = wf_provider
		if self.asign_job_folder: self.asign_job_paths()
		self.create_folder(self.exec_folder)
		if self.linked_folders: # Reset all symlinks to regenerate them in case os wf template changes
								# TODO adapt also to remote execution
			commented_jobs_paths = self.get_paths_for_commented_jobs() # We need to know which symbolic links are from commented jobs, so their results will be used and we CANNOT delete their symbolic link
			self.create_folder(self.raw_folder)
			files = glob.glob(os.path.join(self.exec_folder, '*')) 
			for file in files:
				path = pathlib.Path(file)
				if not file in commented_jobs_paths and path.is_symlink(): os.remove(file) 
		self.create_file('index_execution', self.exec_folder)
		self.launch_all_jobs()
		self.close_file('index_execution')

	def asign_job_paths(self):
		for job in list(self.commands.values()):
			job.asign_folder(self.folder_name)
		for name, job in self.commands.items():
			job.set_dependencies_path()		
			job.name = re.sub('\\)','', job.name) #Clean function characters on name
		for name in list(self.commands.keys()): #Clean function characters on name
			self.commands[re.sub('\\)','', name)] = self.commands.pop(name)

	def get_paths_for_commented_jobs(self):
		commented_jobs_paths = []
		for job in list(self.commands.values()):
			if job.attrib['done']: commented_jobs_paths.append(job.attrib['exec_folder'])
		return commented_jobs_paths

	def launch_all_jobs(self):
		buffered_jobs = []
		sorted_jobs = self.sort_jobs_by_dependencies()
		#print(sorted_jobs, file=sys.stderr)
		for name, job in sorted_jobs:
			if self.verbose: print(job.inspect(all_attribs= True if self.verbose > 1 else False )) # Show job to check by user
			if job.attrib.get('done') != None: self.active_jobs.append(job.name) 

		self.init_log()
		for name, job in sorted_jobs:
			self.write_file('index_execution', f"{name}\t{job.attrib['exec_folder']}") 
			if self.sleep_time > 0: time.sleep(self.sleep_time) 
			if job.attrib.get('done'):
				continue
			else:
				self.rm_done_dependencies(job)
			buffered_jobs = self.launch_job_in_folder(job, name, buffered_jobs)

	def sort_jobs_by_dependencies(self): # We need job ids from queue system so we ask for each job and we give the previous queue system ids as dependencies if necessary
		ar_jobs =[ [name, job] for name, job in self.commands.items() ]
		sorted_jobs = []
		jobs_without_dep = list(filter(lambda job: len(job[1].dependencies) == 0, ar_jobs))
		sorted_jobs = sorted_jobs + jobs_without_dep
		while len(ar_jobs) != len(sorted_jobs):
			ids = [ job[0] for job in sorted_jobs ]
			for job in ar_jobs:
				if job not in sorted_jobs: 
					deps = list(set(job[1].dependencies) - set(ids))
					if len(deps) == 0 : sorted_jobs.append(job)
		return sorted_jobs

	def rm_done_dependencies(self, job):
		remove=[]
		for dependency in job.dependencies:			
			if self.commands[dependency].attrib['done']: remove.append(dependency)
		for rm in remove: job.dependencies.remove(rm)

	def get_relations_and_folders(self):
		relations = {}
		for name, job in self.commands.items(): relations[name] = [job.attrib['exec_folder'], job.dependencies]
		return relations

	def launch_job_in_folder(self, job, id, buffered_jobs):
		if self.linked_folders:
			raw_name = hashlib.sha256(str.encode(job.name)).hexdigest()
			raw_job_folder = os.path.join(self.raw_folder, raw_name)
			self.create_folder(raw_job_folder)
			os.symlink(raw_job_folder, job.attrib['exec_folder']) 
		else:
			self.create_folder(job.attrib['exec_folder'])
		if not job.attrib['buffer']:  # Launch with queue_system the job and all buffered jobs
			self.launch2queue_system(job, id, buffered_jobs)
			buffered_jobs = []#Clean buffer
		else: # Buffer job
			buffered_jobs.append([id, job])
		return buffered_jobs

	def launch2queue_system(self, job, id, buffered_jobs):
		sh_name = job.name+'.sh'
		if self.write_sh:
			self.create_file(sh_name, job.attrib['exec_folder'])
			self.write_file(sh_name, '#!/usr/bin/env bash')
			self.write_file(sh_name, '##JOB_GROUP_ID='+ self.job_identifier)
			self.write_header(id, job, sh_name)
			if self.wf_provider != None:
				virt_command = self.wf_provider.get_virt_init_command(job.attrib['virt_type'], job.attrib['virt'])
				if virt_command != None: self.write_file(sh_name, virt_command)

		#Get dependencies
		#------------------------------------
		ar_dependencies = self.get_dependencies(job, id)
		for id_buff_job, buff_job in buffered_jobs:
			ar_dependencies = ar_dependencies + self.get_dependencies(buff_job, id_buff_job)
		ar_dependencies = list(set(ar_dependencies))

		if self.write_sh: #Write sh body
			self.write_file(sh_name, 'hostname')
			log_file_path = '/'.join([self.exec_folder, '.wf_log', os.path.basename(job.attrib['exec_folder'])])
			self.write_file(sh_name, f"flow_logger -e {log_file_path} -s {job.name}")
			for id_buff_job, buff_job in buffered_jobs:
				self.write_job(buff_job, sh_name)
				buff_job.attrib['exec_folder'] = job.attrib['exec_folder']

			self.write_job(job, sh_name)
			self.write_file(sh_name, f"flow_logger -e {log_file_path} -f {job.name}")
			self.write_file(sh_name, "echo 'General time'")
			self.write_file(sh_name, "times")
			self.close_file(sh_name, 0o755) 

		queue_id = self.submit_job(job, ar_dependencies)
		if not self.verbose: #Append submission data to job
			job.queue_id = queue_id # Returns id of running tag on queue system 
			self.asign_queue_id(buffered_jobs, queue_id)

	def get_dependencies(self, job, id = None):
		ar_dependencies = job.dependencies
		if id != None: 
			if id in ar_dependencies: ar_dependencies.remove(id) #Delete autodependency
		return ar_dependencies

	def asign_queue_id(self, ar_jobs, id):
		for id_job, job in ar_jobs: job.queue_id = id

	def write_job(self, job, sh_name):
		if job.initialization != None: self.write_file(sh_name, job.initialization) 
		if self.comment:
			cmd = '#' + job.parameters
		else:
			if self.extended_logging:
				log_command = '/usr/bin/time -o process_data -v '
			else:
				log_command = 'time '
			cmd = log_command + job.parameters
		self.write_file(sh_name, cmd)

	# Task logging
	# ------------------------------------------------------------------------
	def init_log(self): #TODO adapt to remote execution
		log_path = ('/').join([self.exec_folder, '.wf_log']) #Join must assume linux systems so os.path.join canot be used for windows hosts
		log = parse_log(log_path)
		job_relations_with_folders = self.get_relations_and_folders()
		if self.write_sh:
			self.create_file('wf.json', self.exec_folder)
			self.write_file('wf.json', json.dumps(job_relations_with_folders))
			self.close_file('wf.json')

		for task in self.active_jobs:
			query = log.get(task)
			current_time = int(time.time())
			if query == None:
				log[task] = {'set': [current_time]}
			else:
				log[task]['set'].append(current_time)
		write_log(log, log_path, job_relations_with_folders, write_task_logs= not self.verbose) 
