import os
import re
import sys
import tempfile
from snakemake.utils import logger, min_version, update_config

from default_values import *

# get default values and update them with values specified in config file
default_config = make_default_config()
update_config(default_config, config)
config = default_config


# required snakemake version
min_version("4.3")
shell.executable("/bin/bash")
shell.prefix("set -o pipefail; ")

def get_conda_envs_dir():
    if config.get("yaml_dir"):
        yaml_dir = config.get("yaml_dir")
    else:
        yaml_dir = os.path.join(os.path.dirname(os.path.abspath(workflow.snakefile)), "envs")
    if not os.path.exists(yaml_dir):
        sys.exit("Unable to locate the environmental dependencies file; tried %s" % yaml_dir)
    return yaml_dir


def get_temp_dir(config):
    if config.get("tmpdir"):
        tmp_dir = config["tmpdir"]
    else:
        tmp_dir = tempfile.gettempdir()
    return tmp_dir


def get_bin_summary_files(do_binning, samples):
    ret_str = ""
    if do_binning:
        ret_str = expand("{sample}/genomic_bins/{sample}.summary",
                      sample=samples)
    return ret_str


def get_shell_prefix(config, override={}):
    pfx = config.get("prefix")
    if not pfx:
        return ""

    keys = re.findall(r"__(\w+)__", pfx)
    for k in keys:
        if k in override:
            v = override[k]
        else:
            v = config.get(k, "")
        pfx = pfx.replace("__{key}__".format(key=k), str(v))
    return pfx


def update_config_file_paths(config):
    for sample in config["samples"]:
        try:
            # convert string into list
            if isinstance(config["samples"][sample]["fastq"], str):
                config["samples"][sample]["fastq"] = [config["samples"][sample]["fastq"]]
        # fastq is not required for annotation alone
        except KeyError:
            continue
    return config


def get_quality_controlled_reads(wildcards):
    """Gets quality controlled reads for two cases. When preprocessed with
    ATLAS, returns R1, R2 and se fastq files or just se. When preprocessed
    externaly and run ATLAS workflow assembly, we expect R1, R2 or se.
    """
    n_files = len(config["samples"][wildcards.sample]["fastq"])

    if config.get("workflow", "complete") == "annotate":
        # QA'd reads; the user wants to begin at assembly step
        if n_files == 2:
            fastq = dict(zip(['R1','R2'], config["samples"][wildcards.sample]["fastq"]))
        elif n_files == 1:
            if config["samples"][wildcards.sample].get("paired", False):
                logger.critical(("This protocol requires interleaved reads be "
                                 "reformatted into separate R1 and R2 files "
                                 "before execution. We recommend using "
                                 "reformat.sh from bbtools to convert."))
                sys.exit(1)
            fastq = {'se': config["samples"][wildcards.sample]["fastq"]}
    else:
        # reads that have gone through ATLAS QC
        fractions = ['R1', 'R2', 'se'] if n_files == 2 \
                        or config["samples"][wildcards.sample].get("paired", False) \
                        else ['se']
        fastq = dict(zip(fractions, expand("{sample}/sequence_quality_control/{sample}_QC_{fraction}.fastq.gz",
                                           fraction=fractions, **wildcards)))
    return fastq


def input_params_for_bbwrap(wildcards, input):
    """This function generates the input flag needed for bbwrap for all cases
    possible for get_quality_controlled_reads.
    """
    if hasattr(input, 'R1') and hasattr(input, 'R2'):
        if hasattr(input, 'se'):
            flag = "in1={R1},{se} in2={R2},null".format(**input)
        else:
            flag = "in1={R1} in2={R2}".format(**input)
    elif hasattr(input, 'se'):
        flag = "in1={se}".format(**input)
    else:
        logger.critical(("File input expectation is one of: "
                         "1 file = single-end, "
                         "2 files = R1,R2, or"
                         "3 files = R1,R2,se"
                         "got: {n} files:\n{}").format('\n'.join(input),
                                                       n=len(input)))
        sys.exit(1)
    return flag


CONDAENV = get_conda_envs_dir()

localrules: all
if config.get("workflow") != "download":
    config = update_config_file_paths(config)
    TMPDIR = get_temp_dir(config)
    CHECKMDIR = os.path.join(os.path.dirname(os.path.realpath(config.get("diamond_db", "."))), "checkm")
    SAMPLES = [i for i in config["samples"].keys()]

# quality control of sequences only
if config.get("workflow", "complete") == "qc":
    # define paired end or not
    PAIRED_END = all([config["samples"][s].get("paired", False) \
                     or (len(config["samples"][s]["fastq"]) > 1) for s in SAMPLES])
    MULTIFILE_FRACTIONS = ['R1', 'R2', 'se'] if PAIRED_END else ['se']
    RAW_INPUT_FRACTIONS = ['R1', 'R2'] if PAIRED_END else ['se']

    rule all:
        input:
            expand("{sample}/sequence_quality_control/finished_QC",sample=SAMPLES),
            "finished_QC"

    include: "rules/qc.snakefile"

# quality control, assembly, annotation, and quantification
elif config.get("workflow", "complete") == "complete":
    # define paired end or not
    PAIRED_END = all([config["samples"][s].get("paired", False) \
                     or (len(config["samples"][s]["fastq"]) > 1) for s in SAMPLES])
    MULTIFILE_FRACTIONS = ['R1', 'R2', 'se'] if PAIRED_END else ['se']
    RAW_INPUT_FRACTIONS = ['R1', 'R2'] if PAIRED_END else ['se']

    rule all:
        input:
            expand("{sample}/sequence_quality_control/finished_QC",sample=SAMPLES),
            "finished_QC",
            expand("{sample}/assembly/contig_stats/prefilter_contig_stats.txt",
                sample=SAMPLES),
            expand("{sample}/assembly/contig_stats/final_contig_stats.txt",
                sample=SAMPLES),
#            get_bin_summary_files(config.get("perform_genome_binning", True), SAMPLES),  # this produces an error when genome binning is not active
            expand("{sample}/{sample}_annotations.txt", sample=SAMPLES)

    include: "rules/assemble.snakefile"
    include: "rules/qc.snakefile"

elif config.get("workflow") == "download":

    DBDIR = os.path.realpath(config["db_dir"])
    CHECKMDIR = os.path.join(DBDIR, "checkm")
    CHECKM_ARCHIVE = "checkm_data_v1.0.9.tar.gz"
    FILES = {"adapters.fa": "ae839dc79cfb855a1b750a0d593fe01e",
             "phiX174_virus.fa": "82516880142e8c89b466bc6118696c47",
             "refseq.db": "42b8976656f2cfd661b8a299d6e24c19",
             "refseq.dmnd": "c01facc7e397270ccb796ea799a09108",
             "refseq.tree": "469fcbeb15dd0d4bf8f1677682bde157",
             "silva_rfam_all_rRNAs.fa": "f102e35d9f48eabeb0efe9058559bc66",
             CHECKM_ARCHIVE: "631012fa598c43fdeb88c619ad282c4d"}

    rule all:
        input:
            expand("{dir}/{filename}", dir=DBDIR, filename=list(FILES.keys())),
            "%s/taxon_marker_sets.tsv" % CHECKMDIR

    include: "rules/download.snakefile"

elif config.get("workflow") == "annotate":
    rule all:
        input:
            expand("{sample}_annotations.txt", sample=SAMPLES),
            expand("{sample}/contig_stats.txt", sample=SAMPLES)

    include: "rules/annotate.snakefile"

else:
    print("Workflow %s is not a defined workflow." % config.get("workflow", "[no --workflow specified]"),
          file=sys.stderr)
