import click
import subprocess
import os
import sys
import shutil
from magneto.scripts.get_samples import run_samples
from yaml import load, dump
try:
    from yaml import CLoader as Loader, CDumper as Dumper
except ImportError:
    from yaml import Loader, Dumper

src = os.path.dirname(__file__)


def init_workflow(workflow, skip_qc, merge, dummy, wd):
    """
        redirect target option to targeted rules names
    """

    f_CONFIG = open(os.path.join(wd, "config", "config.yaml"), "r")
    config = load(f_CONFIG, Loader=Loader)
    f_CONFIG.close()


    target= ""
    if workflow == "qc":
        target = "run_QC"
        if skip_qc or config["skip_QC"]:
            raise ValueError("QC require QC ! \nremove skip-qc flag to run QC")
    elif workflow == "motus":
        target = "run_motus"
    elif workflow == "assembly":
        target = "run_assembly"
    elif workflow == "simka":
        target = "target_simka"
    elif workflow == "single_assembly":
        target = "run_assembly"
    elif workflow == "co-assembly":
        target = "run_assembly"
    elif workflow == "binning":
        target = "run_binning"
    elif workflow == "genes":
        target = "run_gene_cat"
    elif workflow == "DM":
        target = "aggregate_deepmased"
    elif workflow == "checkm":
        target = "target_checkm_statistics"
    elif workflow == "dereplicate":
        target = "expand_target_dereplication"
    elif workflow == "none":
        target = ""
    elif workflow == "genomes":
        target = "run_genomes"
    return target

def get_snakefile(file="Snakefile"):
    sf = os.path.join(os.path.dirname(os.path.abspath(__file__)), file)
    if not os.path.exists(sf):
        sys.exit(f"Unable to locate the Snakemake workflow file; tried {sf}")
    return sf


@click.group()
def cli():
    """
        pipeline description here
    """

#####################
#    init command   #
#####################

@click.command()
#@click.argument("samples")
@click.option("--wd", default=os.path.join("."))
def init(wd):
    # init config and sample files
    if not os.path.isdir(wd):
        os.mkdir(wd)
    if not os.path.isdir(os.path.join(wd, "config")):
        os.mkdir(os.path.join(wd, "config"))
    if not os.path.isdir(os.path.join(wd, "logs")):
        os.mkdir(os.path.join(wd, "logs"))
    if not os.path.isdir(os.path.join(wd, "logs/cluster_logs")):
        os.mkdir(os.path.join(wd, "logs/cluster_logs"))
    if not os.path.isdir(os.path.join(wd, "scripts")):
        os.mkdir(os.path.join(wd, "scripts"))

    source_sh_file = os.path.join(src, "scripts", "Fasta_to_Contig2Bin.sh") 
    destination_sh_file = os.path.join(wd, "scripts", "Fasta_to_Contig2Bin.sh")
    shutil.copy(source_sh_file, destination_sh_file)

    #run_samples(samples,os.path.join(wd, "config"))
    f_content = open(os.path.join(src, ".configtemplate.yaml")).read()
    f_content = f_content.replace('name : ""', f'name : {os.path.basename(wd)}')
    f_content = f_content.replace('project : ""', f'project : {os.path.abspath(wd)}')
    f_content = f_content.replace('samples : ""', f'samples : {os.path.abspath(wd)}/config/samples.yaml')
    if not os.path.isdir(os.path.join(wd, "config")):
        os.mkdir(os.path.join(wd, "config"))
    f_CONFIG = open(os.path.join(wd, "config" ,"config.yaml"), 'w')
    f_CONFIG.write(f_content)
    f_CONFIG.close()
    #init sge profile
    if not os.path.isdir(os.path.join(wd, "config", "sge")):
        os.mkdir(os.path.join(wd, "config", "sge"))
    if not os.path.isdir(os.path.join(wd, "config", "slurm")):
        os.mkdir(os.path.join(wd, "config", "slurm"))
    cluster_sge_template = os.path.join(src, ".clustertemplate.sge.yaml")
    cluster_slurm_template = os.path.join(src, ".clustertemplate.slurm.yaml")
    samples_template = os.path.join(src, ".dummy_samples_template.yaml")
    sge_profile = os.path.join(src, ".sge", "config.yaml")
    slurm_profile = os.path.join(src, ".slurm", "config.yaml")
    os.system('cp {} {}'.format(cluster_slurm_template, os.path.join(wd, "config", "cluster_slurm.yaml")))
    os.system('cp {} {}'.format(cluster_sge_template, os.path.join(wd, "config", "cluster_sge.yaml")))
    os.system('cp {} {}'.format(sge_profile, os.path.join(wd, "config", "sge", "config.yaml")))
    os.system('cp {} {}'.format(samples_template, os.path.join(wd, "config/dummy_samples.yaml")))
    os.system(f'cp {sge_profile} {wd}/config/sge/config.yaml')
    os.system(f'cp {slurm_profile} {wd}/config/slurm/config.yaml')


#####################
#   run command     #
#####################

@click.command(context_settings=dict(ignore_unknown_options=True, ))
@click.option("-v", "--verbose", is_flag=True, help="Enables verbose mode")
@click.argument("workflow",
                default = "all",
                type = click.Choice([
                "all",
                "qc",
                "motus",
                "simka",
                "assembly",
                "single_assembly",
                "co-assembly",
                "binning",
                "genes",
                "checkm",
                "dereplicate",
                "genomes",
                "none"])
                )

@click.option("--skip-qc", is_flag=True)
@click.option("--merge", is_flag=True)
@click.option("--dummy", is_flag=True)
@click.option("--wd", default=os.path.join("."))

@click.argument("snakargs", nargs=-1, type=click.UNPROCESSED)

def run(verbose, workflow, skip_qc, merge, dummy, wd, snakargs):

    target = init_workflow(workflow, skip_qc, merge, dummy, wd)

    params = ["--config"]
    params.append("target={}".format(workflow))
    if skip_qc:
        params.append("skip_QC=True")
    if merge:
        params.append("merging=True")
    if dummy:
        params.append("samples=config/dummy_samples.yaml")

    default_snakargs = ["--printshellcmds"]#,"--stats ./reports/workflow_stats.json"]
    args_list = params + list(snakargs) + default_snakargs

    print(" ".join(args_list))

    if workflow == "":
        exit(1)
    else:
        pass

    cmd = (
        "snakemake --snakefile {snakefile} "
        "--configfile '{config_file}' "
        "--directory '{working_dir}' "
        "--use-conda "
        "{target_rule} "
        #"--cores 2 " #only for local test
        "{snakemake_args} "
        ).format(
        snakefile      = get_snakefile(),
        config_file    = os.path.join(wd, "config", "config.yaml"),
        working_dir    = wd,
        target_rule    = target,
        snakemake_args = " ".join(args_list)
    )

    subprocess.check_call(cmd, shell=True)

cli.add_command(init)
cli.add_command(run)

if __name__ == '__main__':
    cli()
