import sys, os, re
import pytest
from py_autoflow.stack import Stack

ROOT_PATH=os.path.dirname(__file__)
TEMPLATES_PATH = os.path.join(ROOT_PATH, 'templates')

@pytest.fixture()
def stack():
	stack = Stack('exec', {
		'cpu' : 16, 
		'mem' : '4gb', 
		'time' : '20:00:00',
		'node' : None,
		'multinode' : 0,
		'ntask' : False,
		'key_name' : False,
		'retry' : False,
		'Variables' : None,
		'workflow' : '',
		'white_list': [],
		'black_list': []
	})
	return(stack)

# test scan_nodes method
#---------------------------------------------

def test_scan_simple_node(stack):
	node_lines = [
		#Single node
		"result){\n",
			"null\n", 
		"?\n", 
			"touch algo\n", 
	  	"}\n"
	]
	
	result = [
		["result)", "null\n", "touch algo\n"]
	]
	
	test = stack.scan_nodes(node_lines)
	assert(result, test)

def test_scan_double_node(stack):
	node_lines = [
		#Single node
		"algo){\n", 
			"null\n", 
		"?\n", 
			"echo 'OK'\n", 
	  	"}\n",
	  	#Single node 
	  	"result){\n", 
			"null\n", 
		"?\n", 
			"touch algo\n", 
		"}\n"
	]

	result = [
	  ["algo)", "null\n", "echo 'OK'\n"], 
	  ["result)", "null\n", "touch algo\n"]
	]

	test = stack.scan_nodes(node_lines)
	assert(result, test)

def test_scan_iterative_node(stack):
	node_lines = [
	  #Nested iterative nodes
	  "itera_[11;22]){\n", 
		"algo_[aa;bb]){\n", 
		  "null\n", 
		  "?\n", 
		  "echo 'OK'itera_(+)(*)\n", 
		"}\n", 
	  "}\n"
	]

	result = [
	  ["itera_[11;22])", "", [1]], 
	  ["algo_[aa;bb])", "null\n", "echo 'OK'itera_(+)(*)\n"]
	]

	test = stack.scan_nodes(node_lines)
	assert(result, test)

def test_scan_3_nested_iterative_node(stack):
	node_lines = [
	  #Nested iterative nodes
	  "itera_[11;22]){\n", 
		"?\n", 
		"algo_[aa;bb]){\n", 
		 "?\n", 
		  "mas_[ZZ;YY]){\n", 
			"null\n", 
			"?\n", 
			"echo 'OK'itera_(+)(*)mas_(+)\n", 
		  "}\n", 
		"}\n", 
	  "}\n"
	]

	result = [
	  ["itera_[11;22])", "", [1]], 
	  ["algo_[aa;bb])", "", [2]], 
	  ["mas_[ZZ;YY])", "null\n", "echo 'OK'itera_(+)(*)mas_(+)\n"]
	]

	test = stack.scan_nodes(node_lines)
	assert(result, test)

def test_scan_nested_iterative_node_and_external_single_node(stack):
	#test nested iterative nodes with a final single node (bug)
	node_lines = [
	  #Nested iterative nodes
	  "itera_[11;22]){\n", 
		"algo_[aa;bb]){\n", 
		  "null\n", 
		  "?\n", 
		  "echo 'OK'itera_(+)(*)\n", 
		"}\n", 
	  "}\n", 
	  #Single node
	  "result){\n", 
		"null\n", 
		"?\n", 
		"touch algo\n", 
	  "}\n"
	]

	result = [
	  ["itera_[11;22])", "", [1]], 
	  ["algo_[aa;bb])", "null\n", "echo 'OK'itera_(+)(*)\n"], 
	  ["result)", "null\n", "touch algo\n"]
	] 
	
	test = stack.scan_nodes(node_lines)
	assert(result, test)

# test parse method
#---------------------------------------------
def test_parse_workflow(stack):
	workflow = """algo){ 
		null 
		? 
		echo 'OK'
	  }
	  
	  result){ 
		null 
		?
		touch algo 
	  }"""
	
	stack.workflow = workflow
	stack.parse()
	commands = stack.commands
	algoBt = commands['algo']
	assert(algoBt.main_command, "echo 'OK'\n")
	assert(algoBt.initialization, "null\n")
	assert(algoBt.iterator, [None])
	assert(algoBt.parent, None)
	assert(algoBt.attrib, {'cpu': 1, 'mem': None, 'time': None, 'node': None, 'multinode': None, 'ntask': None, 'additional_job_options': None, 'done': False, 'folder': True, 'buffer': False, 'exec_folder': 'exec', 'cpu_asign': 'mono'})
	resBt = commands['result']
	assert(resBt.main_command, "touch algo\n")
	assert(resBt.initialization, "null\n")
	assert(resBt.iterator, [None])
	assert(resBt.parent, None)
	assert(resBt.attrib, {'cpu': 1, 'mem': None, 'time': None, 'node': None, 'multinode': None, 'ntask': None, 'additional_job_options': None, 'done': False, 'folder': True, 'buffer': False, 'exec_folder': 'exec', 'cpu_asign': 'mono'})

# test get_jobs_relations method
#---------------------------------------------
def test_parse_workflow(stack):
	workflow = """algo){ 
		null 
		? 
		echo 'OK'
	  }
	  
	  result){ 
		null 
		?
		touch algo 
		echo algo)/file
	  }"""
	
	stack.workflow = workflow
	stack.parse()
	job_relations = stack.get_jobs_relations()
	list_job_relations = []
	for j_name, job in job_relations.items():
		list_job_relations.append([j_name, job.name, job.dependencies])
	true_relations = [
		['algo', 'algo', []],
		['result', 'result', ['algo']]
	]
	assert(true_relations, list_job_relations)

def test_resources_by_file(stack):
	#resources: -s -n bigmem -c 1 -t "7-00:00:00" -m "300gb"
	workflow = """algo){ 
		resources: -r test
		? 
		echo 'OK'
		echo [cpu]
	  }
	"""
	stack.parse_resources({'resources': os.path.join(TEMPLATES_PATH, 'resources.json')})
	stack.workflow = workflow
	stack.parse()
	job_relations = stack.get_jobs_relations()
	expected = {'cpu': 2, 'mem': '300GB', 'time': '7-00:00:00', 'node': 'bigmem', 'multinode': 0, 'ntask': False, 'additional_job_options': None, 'done': False, 'folder': True, 'buffer': False, 'exec_folder': 'exec', 'cpu_asign': 'number', 'virt': 'test_virt', 'virt_type' : 'env'}
	assert(job_relations['algo'].attrib == expected)

def test_virtualization_parsing_by_file(tmp_dir, stack):
	#resources: -s -n bigmem -c 1 -t "7-00:00:00" -m "300gb"
	output = os.path.join(tmp_dir, 'virtualization')	
	stack.parse_resources({'resources': os.path.join(TEMPLATES_PATH, 'resources_virt.json'), 'output': output})
	virts = stack.wf_manager.virtualizations
	virts['env']['test_virt'] = re.sub(output + '/', '', virts['env']['test_virt'])
	expected = {'env': {'test_virt': '.autoflow_smith/env/test_virt'}, 'sing': {}}
	assert(virts == expected)