# == Native Modules ==
import subprocess
from os.path import abspath
import sys
# == Installed Modules ==
import yaml
# == Project Modules ==
from prog.medit_lib import (compress_file,
                            file_exists,
                            launch_shell_cmd,
                            parse_editor_request,
                            project_file_path,
                            prRed,
                            set_export,
                            write_yaml_to_file
                            )


def guide_prediction(args, jobtag):
	# == Load Run Parameters values ==
	query_input = args.query_input
	user_jobtag = args.user_jobtag
	root_dir = abspath(args.output)
	db_path_full = f"{abspath(args.db_path)}/medit_database"
	mode = args.mode
	private_genomes = args.private_genome
	qtype = args.qtype_request
	editor_request = parse_editor_request(args.editor_request)
	be_request = parse_editor_request(args.be_request)

	# == Load SLURM-related values ==
	ncores = args.ncores
	maxtime = args.maxtime
	parallel_processes = args.parallel_processes
	dry_run = args.dry_run

	# == Define dynamic SMK call variables ==
	allowed_rules = ['']
	cluster_smk_setup = ['']
	smk_verbosity = [True]
	smk_run_triggers = ''
	dryrun_setup = ''

	# ->=== INPUT CHECKS ===<-
	#   == Check the presence of private genome among the inputs ==
	if private_genomes:
		mode = 'private'
		private_genomes = args.private_genome.split(",")
	#   == Check the request to run on a cluster
	if parallel_processes:
		# --> Upon SLURM run request, the guide_prediction.smk is split in two separate runs
		# --> That's because samtool's conda package crashes on a libcrypto error when
		#       it's deployed by snakemake on a SLURM node
		cluster_smk_setup = ['', '--cluster "sbatch -t {cluster.time} -n {cluster.cores}" '
		                         '--cluster-config config/medit_cluster.yaml']
		allowed_rules = ['--until "consensus_fasta"', '']
		smk_verbosity = [False, True]
	#   == Check the dry run request
	if dry_run:
		dryrun_setup = '-n'
	if user_jobtag:
		smk_run_triggers = '--rerun-triggers "mtime"'

	# ->=== OUTPUT SETUP ===<-
	# == Set import paths tied to the SMK pipeline
	config_db_path = f"{db_path_full}/config_db/config_db.yaml"
	# == Set export paths tied to the SMK pipeline ==
	config_dir_path = f"{root_dir}/config"
	# == Set export paths for dynamic YAML files ==
	dynamic_config_path = f"{config_dir_path}/config_{jobtag}.yaml"
	dynamic_cluster_path = f"{config_dir_path}/cluster_{jobtag}.yaml"
	# == Create sub-folders to host VCFs, and config files ==
	set_export(config_dir_path)

	# ->=== CONFIG FILES IMPORT ===<-
	#   == Load template configuration files ==
	config_template_path = project_file_path("smk.config", "medit_guide_pred.yaml")
	cluster_template_path = project_file_path("smk.config", "medit_cluster.yaml")
	with open(config_template_path, 'r') as config_handle:
		config_template = yaml.safe_load(config_handle)
	with open(cluster_template_path, 'r') as cluster_handle:
		cluster_template = yaml.safe_load(cluster_handle)

	#   == Load Database configuration file generated by db_set ==
	try:
		with open(config_db_path, 'r') as config_db_handle:
			config_db = yaml.safe_load(config_db_handle)
	except FileNotFoundError:
		prRed(f"Couldn't fund the file {config_db_path}. "
		      f"Please double-check the path to <medit_database> and provide the path via '-d' option")

	# === Assign Variables to Configuration File ===
	config_template['run_name'] = f"{mode}_{jobtag}"
	config_template['support_tables'] = db_path_full
	config_template['processing_mode'] = mode
	config_template['output_directory'] = root_dir
	config_template['variant_query_dir'] = query_input
	# Assign run parameters to config
	config_template['qtype'] = qtype
	config_template['editor_request'] = editor_request
	config_template['be_request'] = be_request
	# Assign cluster options
	cluster_template['__default__']['cores'] = ncores
	cluster_template['__default__']['time'] = maxtime

	# ->=== PRIVATE GENOME RUN ===<-
	# == Check run mode ==
	if mode == 'private':
		# == Enforce presence of private genome in this mode ==
		if not private_genomes:
			print("Please provide a VCF input file to run mEdit's private mode")
			sys.exit(1)
		# == Create a private VCF directory when a private run is issued
		vcf_dir_path = f"{config_db['meditdb_path']}/{mode}/source_vcfs"
		set_export(vcf_dir_path)
		# == Check the config_db file for instructions on which human genome version should be used
		latest_reference = bool(config_db["latest_reference"])
		custom_reference = bool(config_db["custom_reference"])
		if latest_reference:
			config_template['sequence_id'] = ["latest_hg38"]
		if custom_reference:
			config_template['sequence_id'] = ["custom_reference"]
		# == VCF ID adjustment for private vcf run ==
		#   => Import VCF file prefix information to config file
		tagged_genomes = []
		count_tag = 1
		for private_genome in private_genomes:
			loop_tag = f"{jobtag}_{count_tag}"
			vcf_filename = f"{loop_tag}.vcf"
			tagged_genomes.append(loop_tag)
			#   => Avoid re-creating private VCF mirrors within the medit DB
			if not file_exists(f"{vcf_dir_path}/{vcf_filename}"):
				if not file_exists(f"{vcf_dir_path}/{vcf_filename}.gz"):
					#   => Create a copy of the VCF in the internal mEdit directory
					launch_shell_cmd(f"cp {private_genome} {vcf_dir_path}/{vcf_filename}", True)
					#   => Check VCF file compression and compress if necessary
					compress_file(f"{vcf_dir_path}/{vcf_filename}")
			count_tag += 1
		#   => Add any amount of private genomes to the config file
		config_template["vcf_id"] = tagged_genomes

	# === Write YAML configs to mEdit Root Directory ===
	write_yaml_to_file(config_template, dynamic_config_path)
	write_yaml_to_file(cluster_template, dynamic_cluster_path)

	# === Invoke SMK Pipelines ===

	print("# == Calling Guide Prediction pipeline == #")
	for smk_setup_idx in range(len(allowed_rules)):
		try:
			# --> When cluster submission is switched on,
			launch_shell_cmd(f"snakemake "
			                 f"--snakefile {project_file_path('smk.pipelines', 'guide_prediction.smk')} "
			                 # f"--directory {project_file_path('smk.pipelines', 'guide_prediction.smk')} "
			                 f"-j {ncores} "
			                 f"{smk_run_triggers} "
			                 f"{allowed_rules[smk_setup_idx]} "
			                 f"{cluster_smk_setup[smk_setup_idx]} "
			                 f"--configfile {config_db_path} "
			                 f"{dynamic_config_path} "
			                 f"--use-conda "
			                 f"{dryrun_setup}",
			                 smk_verbosity[smk_setup_idx]
			                 )
		except subprocess.CalledProcessError as e:
			print(f"Error: {e}")
		except ValueError:
			print(f"Process completed in a previous run. Moving to the next one...")
