# == Native Modules ==
from copy import copy
from os.path import abspath
from pathlib import Path
import pickle
import subprocess
import sys
# == Installed Modules ==
import yaml
# == Project Modules ==
from prog.medit_lib import (compress_file,
							check_format,
							file_exists,
							handle_shell_exception,
							launch_shell_cmd,
							project_file_path,
							prRed,
							set_export,
							validate_editor_list,
							write_yaml_to_file)


def guide_prediction(args, jobtag):
	# === Load Run Parameters values ===
	# == Inputs and File tree setup
	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 and associated values
	mode = args.mode
	mode_name = copy(mode)
	custom_vcfs = args.custom_vcf
	qtype = args.qtype_request
	# == Editor request and associated values
	editor_request = args.editor_request
	be_request = args.be_request
	cutdist = args.cutdist
	# == Custom editor required values
	pam = args.pam
	guide_length = args.guide_length
	pam_is_first = args.pam_is_first
	target_base = args.target_base
	result_base = args.result_base
	dsb_position = args.dsb_position
	editing_window = args.editing_window

	# === 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 = ''

	# ->=== 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 Input Query Files
	query_input_path = f"{root_dir}/queries"
	# == 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:
		handle_shell_exception("ConfigNotFound", config_db_path, True)
		exit(0)

	# ->=== INPUT CHECKS ===<-
	# == Check the input paths and copy them over to mEdits Filetree
	query_input_list = query_input.split(',')
	formatted_query_input_list = []
	query_index_count = 0
	# == Check file integrity
	for query_filename in query_input_list:
		if not file_exists(query_filename):
			handle_shell_exception(str('FileNotFound'), str(query_filename), True)
			exit(0)
		# == Copy query files in the right naming format for mEdit
		if mode_name == 'fast':
			mode_name = 'standard'
		formatted_query_filename = f"{query_input_path}/{mode_name}_{jobtag}_{query_index_count}.csv"
		formatted_query_input_list.append(formatted_query_filename)
		query_index_count += 1
		set_export(query_input_path)
		#   => Create a copy of the VCF in the internal mEdit directory
		launch_shell_cmd(f"cp {query_filename} {formatted_query_filename}", True)
	#   == Check the presence of custom vcf genome among the inputs ==
	if custom_vcfs:
		mode = 'vcf'
		custom_vcfs = args.custom_vcf.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"'
	# == Check lists of editors and validate entries
	if editor_request not in ['clinical', 'custom']:
		with open(str(config_db["editors"]), 'rb') as editors_pkl_handle:
			editors_dict = pickle.load(editors_pkl_handle)
		editor_request = validate_editor_list(editor_request, list(editors_dict['all'].keys()), '--editor')
	if be_request not in ['default', 'custom']:
		with open(str(config_db["base_editors"]), 'rb') as be_pkl_handle:
			base_editors_dict = pickle.load(be_pkl_handle)
		be_request = validate_editor_list(be_request, list(base_editors_dict['all'].keys()), '--be')

	# == Check custom editors parameters
	if editor_request == 'custom':
		pam = check_format(pam, str, 'pam', 'XXX')
		pam_is_first = check_format(pam_is_first, bool, 'pamisfirst', None)
		guide_length = check_format(guide_length, int, 'guidelen', -1)
		dsb_position = check_format(dsb_position, int, 'dsb_pos', -10000)
	elif be_request == 'custom':
		pam = check_format(pam, str, 'pam', 'XXX')
		pam_is_first = check_format(pam_is_first, bool, 'pamisfirst', None)
		guide_length = check_format(guide_length, int, 'guidelen', -1)
		editing_window = check_format(editing_window, tuple, 'edit_win', (0, 0))
		target_base = check_format(target_base, str, 'target_base', 'X')
		result_base = check_format(result_base, str, 'result_base', 'X')

	# ->=== CHECK RUN MODE ===<-
	if mode == 'vcf':
		# == Enforce presence of custom vcf genome in this mode ==
		if not custom_vcfs:
			print("Please provide a VCF input file to run mEdit's vcf mode")
			sys.exit(1)
		# == Create a custom VCF directory when a vcf 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 custom vcf run ==
		#   => Import VCF file prefix information to config file
		tagged_genomes = []
		count_tag = 1
		for custom_vcf in custom_vcfs:
			loop_tag = f"{jobtag}_{count_tag}"
			vcf_filename = f"{loop_tag}.vcf"
			tagged_genomes.append(loop_tag)
			#   => Avoid re-creating custom 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 {custom_vcf} {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 custom vcf genomes to the config file
		config_template["vcf_id"] = tagged_genomes
	elif mode == 'fast':
		allowed_rules = ['--until "predict_guides" --omit-from consensus_fasta']
		# == 'fast' mode is a sub-mode of 'standard';
		# downstram processes must acknowledge that this is a standard run
		mode = 'standard'

	# === 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'] = list(formatted_query_input_list)
	config_template['query_index'] = list(range(len(formatted_query_input_list)))
	# Assign run parameters to config
	config_template['qtype'] = qtype
	config_template['editor_request'] = editor_request
	config_template['be_request'] = be_request
	config_template['distance_from_cutsite'] = cutdist
	# Assign custom editor parameters to config
	config_template['pam'] = pam
	config_template['guide_length'] = guide_length
	config_template['pam_is_first'] = pam_is_first
	config_template['dsb_position'] = dsb_position
	config_template['editing_window'] = editing_window
	config_template['target_base'] = target_base
	config_template['result_base'] = result_base
	# Assign cluster options
	cluster_template['__default__']['cores'] = ncores
	cluster_template['__default__']['time'] = maxtime

	# === 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"--keep-going "
							 f"--rerun-incomplete "
							 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...")
