import click
import subprocess
import os
import sys
import shutil
import yaml
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, 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 == "binning":
        target = "run_binning"
    elif workflow == "genes":
        target = "run_gene_cat"
    elif workflow == "none":
        target = ""
    elif workflow == "genomes":
        target = "run_genomes"
    elif workflow == "metaT":
        target = "run_metaT"        
    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('database : ""', f'database : {os.path.abspath(wd)}/Database')
    f_content = f_content.replace('metaT_samples : ""', f'metaT_samples : {os.path.abspath(wd)}/config/metaT_samples.yaml')    
    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")
    metaG_samples_template = os.path.join(src, ".metaG_samples_template.yaml")
    metaT_samples_template = os.path.join(src, ".metaT_samples_template.yaml")    
    sge_profile = os.path.join(src, ".sge", "config.yaml")
    slurm_profile = os.path.join(src, ".slurm", "config.yaml")
    os.system(f'cp {cluster_slurm_template} {wd}/config/cluster_slurm.yaml')
    os.system(f'cp {cluster_sge_template} {wd}/config/cluster_sge.yaml')
    os.system(f'cp {metaG_samples_template} {wd}/config/samples.yaml')
    os.system(f'cp {metaT_samples_template} {wd}/config/metaT_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.option("--skip-qc", is_flag=True)
@click.option("--merge", is_flag=True)
@click.option("--wd", default=os.path.join("."))
@click.option("--strategies", default="SACB", help="List of strategies to activate, separated by commas (e.g., SASB,SACB)")
@click.argument("workflow",
                default = "all",
                type = click.Choice([
                "all",
                "qc",
                "motus",
                "assembly",
                "binning",
                "genes",
                "genomes",
                "metaT",
                "none"])
                )
@click.argument("snakargs", nargs=-1, type=click.UNPROCESSED)
def run(verbose, workflow, skip_qc, merge, wd, strategies, snakargs):
    target = init_workflow(workflow, skip_qc, merge, wd)

    strategy_list = []
    if strategies:
        strategy_list = [s.strip() for s in strategies.split(",")]

        # Check SA*/CA* incompatibility
        sa = [s for s in strategy_list if s.startswith("SA")]
        ca = [s for s in strategy_list if s.startswith("CA")]
        if sa and ca:
            raise click.UsageError(
                f"Incompatible strategy combination: {', '.join(sa)} cannot be used with {', '.join(ca)}.\n"
                "You cannot mix single-assembly (SA*) and co-assembly (CA*) strategies."
            )

        print(f"✔ Strategies selected: {strategy_list}")

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

    if strategy_list:
        strategies_str = ",".join(strategy_list)
        params.append(f"binning_strategies={strategies_str}")

    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()