import datetime
import glob
import gzip
import io
import itertools
import numpy
import os
import pandas as pd
import re
import shutil
import snakemake.utils as smk_utils
import sys
import tarfile
import utils
import yaml

from Bio import SeqIO
from collections import defaultdict
from pytools.persistent_dict import PersistentDict
from utils import PROPHANE_VERSION
from utils.config import validate_config_and_set_defaults, read_general_config
from utils.input_output import (
    get_headline,
    iter_file,
    load_yaml,
    read_taxmap,
    write_yaml)
from utils.db_handling.db_access import DbAccessor, get_db_object_for_dbtype_and_version
from utils.helper_funcs import get_timestamp
from utils.helper_funcs_for_snakemake_workflow import (
    get_db_obj_for_task_and_dbdir,
    get_mafft_report,
    get_task_lca,
    get_task_map,
    get_task_plot_file,
    get_taxblast_result,
    get_hmmer_version_for_db_type, )
from warnings import warn

# explicitly set shell used by snakemake to bash
shell.executable("/bin/bash")

# set minimum snakemake version
smk_utils.min_version("5.7")


if not config:
    sys.exit("error: no config file submitted")

job_config = validate_config_and_set_defaults(
    config, schema=os.path.join(workflow.basedir, "schemas", "config.schema.yaml")
)

# try to validate job config with general config schema
try:
    smk_utils.validate(job_config,
                       schema=os.path.join(workflow.basedir, "schemas", "general_config.schema.yaml"))
except smk_utils.WorkflowError:
    if "use_test_general_config" in job_config:
        general_config_yaml = job_config['use_test_general_config']
        warn("'use_test_general_config' key in job-config is deprecated, please directly include general config params.")
    else:
        general_config_yaml = ""
    general_config = read_general_config(general_config_yaml=general_config_yaml)
    job_config['db_base_dir'] = general_config['db_base_dir']
finally:
    config = job_config

LCATASKS = list(range(len(config['tasks'])))


def trace(msg, die=False):
    '''prints a TRACE message to stdout and optionally stderr which causes program termination'''
    for m in msg.split("\n"):
        print(m)
    if die:
        sys.exit(msg)

def get_all_plots(wildcards):
    plots = []
    db_base_dir = get_db_base_dir()
    for taskid, task_dict in enumerate(config["tasks"]):
        plots.append(get_task_plot_file(taskid, task_dict, db_base_dir))
    return plots

def get_db_base_dir():
    return config['db_base_dir']

def get_db_yaml_for_task(task_config_dict, db_acc=None):
    db_base_dir = config['db_base_dir']
    if not db_acc:
        db_acc = DbAccessor(db_base_dir)
    db_obj = get_db_obj_for_task_and_dbdir(task_config_dict, db_base_dir, db_acc)
    yaml_file = db_obj.get_yaml()
    return yaml_file

def get_db_object_for_task(task_config_dict, db_acc=None):
    db_base_dir = config['db_base_dir']
    if not db_acc:
        db_acc = DbAccessor(db_base_dir)
    db_obj = get_db_obj_for_task_and_dbdir(task_config_dict, db_base_dir, db_acc)
    return db_obj


workdir: config['output_dir']

# set container dir to avoid writing to prophane bin dir
storage = PersistentDict("mystorage", container_dir=config['output_dir'])
onstart:
    print("job:", config['general']['job_name'])
    # global vars within 'onstart' are not persistent, outside of 'onstart' the variable is set multiple times
    # start time is therefore stored in a persistent dict
    storage.store("START_TIME", get_timestamp())


db_setup_config = os.path.join(config["db_base_dir"], "db_setup_config.yaml")
if not os.path.exists(db_setup_config):
    if not os.path.exists(config["db_base_dir"]):
        os.mkdir(config["db_base_dir"])
    write_yaml({"taxlevel": config["taxlevel"]}, db_setup_config)

subworkflow db_setup_workflow:
    workdir:
        config["db_base_dir"]
    snakefile:
        "Snakefile_db_setup.smk"
    configfile: db_setup_config


rule all:
    input:
        'summary.txt',
        get_all_plots,
        'job_info.txt'
    threads: 3


def make_path_relative_to_db_base_dir(f):
    return os.path.relpath(f, config['db_base_dir'])

def get_all_preprocessed_db_files(w):
    files_abs = []
    for task_dct in config['tasks']:
        if task_dct['prog'] == 'acc2annot_mapper':
            continue
        db_obj = get_db_object_for_task(task_dct)
        prog = task_dct['prog']
        files_abs.append(db_obj.get_yaml())
        files_abs += db_obj.get_all_mandatory_files()
        if not (prog == "emapper" and db_obj.get_version().startswith('5')): # eggnog v5 is the only db without map file
            files_abs.append(db_obj.get_db_map_file())
        db_dir = os.path.dirname(db_obj.get_yaml())
        if prog == "diamond blastp":
            files_abs.append(os.path.join(db_dir,
                                          "processed",
                                          "diamond_db.dmnd"))
            files_abs.append(db_obj.get_tax2annot_map_file())
        elif prog in ["hmmscan", "hmmsearch"]:
            hmmer_vers = get_hmmer_version_for_db_type(db_obj.get_type())
            hmm_lib = os.path.join(db_dir,
                                   "processed",
                                   f"hmmlib-{hmmer_vers}")
            files_abs.append(hmm_lib)
            files_abs += [f"{hmm_lib}.{ext}" for ext in ["h3f", "h3i", "h3m", "h3p"]]
        elif prog == "emapper":
            pass
    files_rel_to_db_base_dir = [make_path_relative_to_db_base_dir(f) for f in sorted(set(files_abs))]
    return [db_setup_workflow(f) for f in files_rel_to_db_base_dir]


rule db_setup:
    input:
        get_all_preprocessed_db_files


#input and data mapping handling
include: "rules/analysis/parse_search_result_into_protein_groups_sql_db.smk"
include: "rules/analysis/map_to_seq.smk"
include: "rules/analysis/create_pg_fasta.smk"
include: "rules/analysis/tax_by_taxmap.smk"

#analyses
include: "rules/analysis/run_mafft.smk"
include: "rules/analysis/report_mafft.smk"
include: "rules/analysis/run_acc2annot_mapper.smk"
include: "rules/analysis/run_diamond.smk"
include: "rules/analysis/run_emapper.smk"
include: "rules/analysis/run_hmmer.smk"
include: "rules/analysis/lca_task.smk"
include: "rules/analysis/quant_pg.smk"
include: "rules/analysis/quant_task.smk"
include: "rules/analysis/create_summary.smk"

#plotting
include: "rules/analysis/create_krona.smk"

#info
include: "rules/analysis/create_info.smk"
