import os
import sys
import json

from py_autoflow.queue_managers.slurm2_manager import Slurm2Manager
from py_autoflow.program import Program
from py_autoflow.wflow_smith import WfSmith
from py_autoflow.cli_manager import *

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


def test_detect_slurm2():
	options = {'show_submit_command': False}
	available = Slurm2Manager.isavailable(options)
	assert available == True

def test_write_sh(tmp_dir):
	execution_folder = os.path.join(tmp_dir, 'slurm2')
	job_attribs = {'folder': True, 'exec_folder': execution_folder, 'buffer': False, 'done': False,
					'ntask': False, 'cpu': 1, 'mem': '1GB', 'time': '0-00:00:01' }
	job = Program('test', '', 'echo test', [], job_attribs)
	jobs = {'test': job}
	execution_options = {'key_name': False, 'asign_job_folder': True, 
						'sleep_time': 0, 'verbose': True, 'write_sh': True,
						'identifier': 'test', 'show_submit_command': False}
	manager = Slurm2Manager(execution_folder, execution_options, jobs)
	manager.exec()
	file_data = read_file(os.path.join(execution_folder, 'echo_0000', 'test.sh'), execution_folder + '/')
	test_result = {'content': file_data['content'], 'comments': file_data['comments']}
	expected_result = { 'content': 'time echo test\n',
						'comments': "\n".join(['#!/usr/bin/env bash', 
									f'##JOB_GROUP_ID={execution_options['identifier']}',
									f'#SBATCH --cpus-per-task={job_attribs['cpu']}',
									f'#SBATCH --mem={job_attribs['mem']}',
									f'#SBATCH --time={job_attribs['time']}',
									'#SBATCH --error=job.%J.err',
									'#SBATCH --output=job.%J.out']) + "\n"
						}
	
	assert expected_result == test_result


def test_write_sh_with_virtualization(tmp_dir):
	execution_folder = os.path.join(tmp_dir, 'slurm2_virt')
	job_attribs = {'folder': True, 'exec_folder': execution_folder, 'buffer': False, 'done': False,
					'ntask': False, 'cpu': 1, 'mem': '1GB', 'time': '0-00:00:01', 
					'virt': 'test_virt', 'virt_type': 'env' }
	job = Program('test', '', 'echo test', [], job_attribs)
	jobs = {'test': job}
	execution_options = {'key_name': False, 'asign_job_folder': True, 
						'sleep_time': 0, 'verbose': True, 'write_sh': True,
						'identifier': 'test', 'show_submit_command': False}
	manager = Slurm2Manager(execution_folder, execution_options, jobs)
	wf_provider = WfSmith(['test_path'])
	test_env_path = os.path.join('test_path', 'test_virt')
	wf_provider.virtualizations['env'] = {'test_virt': os.path.join('test_path', 'test_virt')}
	manager.exec(wf_provider = wf_provider)
	file_data = read_file(os.path.join(execution_folder, 'echo_0000', 'test.sh'), execution_folder + '/')
	test_result = {'content': file_data['content'], 'comments': file_data['comments']}
	expected_result = { 'content': f'source {os.path.join(test_env_path, 'bin', 'activate')}\ntime echo test\n',
						'comments': "\n".join(['#!/usr/bin/env bash', 
									f'##JOB_GROUP_ID={execution_options['identifier']}',
									f'#SBATCH --cpus-per-task={job_attribs['cpu']}',
									f'#SBATCH --mem={job_attribs['mem']}',
									f'#SBATCH --time={job_attribs['time']}',
									'#SBATCH --error=job.%J.err',
									'#SBATCH --output=job.%J.out']) + "\n"
						}
	
	assert expected_result == test_result


def test_send_sh(tmp_dir):
	execution_folder = os.path.join(tmp_dir, 'slurm2send')
	job_attribs = {'folder': True, 'exec_folder': execution_folder, 'buffer': False, 'done': False,
					'ntask': False, 'cpu': 1, 'mem': '1GB', 'time': '0-00:05:00', 'node': 'cal'}
	job = Program('test', '', 'echo test', [], job_attribs)
	jobs = {'test': job}
	execution_options = {'key_name': False, 'asign_job_folder': True, 
						'sleep_time': 0, 'verbose': False, 'write_sh': True,
						'identifier': 'test', 'show_submit_command': False, 'external_dependencies': []}
	manager = Slurm2Manager(execution_folder, execution_options, jobs)
	manager.exec()
	
	assert job.queue_id != None
