"""Handle installation and updates of bcbio-nextgen, third party software and data.

Enables automated installation tool and in-place updates to install additional
data and software.
"""
from __future__ import print_function
import argparse
import collections
import contextlib
import datetime
import dateutil
from distutils.version import LooseVersion
import gzip
import os
import shutil
import subprocess
import sys
import glob

import requests
from six.moves import urllib
import toolz as tz
import yaml

from bcbio import broad, utils
from bcbio.pipeline import genome
from bcbio.variation import effects
from bcbio.distributed.transaction import file_transaction
from bcbio.pipeline import datadict as dd

REMOTES = {
    "requirements": "https://raw.githubusercontent.com/chapmanb/bcbio-nextgen/master/requirements-conda.txt",
    "gitrepo": "https://github.com/chapmanb/bcbio-nextgen.git",
    "cloudbiolinux": "https://github.com/chapmanb/cloudbiolinux/archive/master.tar.gz",
    "genome_resources": "https://raw.github.com/chapmanb/bcbio-nextgen/master/config/genomes/%s-resources.yaml",
    "snpeff_dl_url": ("http://downloads.sourceforge.net/project/snpeff/databases/v{snpeff_ver}/"
                      "snpEff_v{snpeff_ver}_{genome}.zip")}
SUPPORTED_GENOMES = ["GRCh37", "hg19", "hg38", "hg38-noalt", "mm10", "mm9",
                     "rn6", "rn5", "canFam3", "dm3", "galGal4", "phix",
                     "pseudomonas_aeruginosa_ucbpp_pa14", "sacCer3", "TAIR10",
                     "WBcel235", "xenTro3", "GRCz10"]
SUPPORTED_INDEXES = ["bowtie", "bowtie2", "bwa", "novoalign", "rtg", "snap",
                     "star","twobit", "seq", "hisat2"]
DEFAULT_INDEXES = ["rtg", "twobit"]

Tool = collections.namedtuple("Tool", ["name", "fname"])

def upgrade_bcbio(args):
    """Perform upgrade of bcbio to latest release, or from GitHub development version.

    Handles bcbio, third party tools and data.
    """
    print("Upgrading bcbio")
    args = add_install_defaults(args)
    if args.upgrade in ["stable", "system", "deps", "development"]:
        if args.upgrade == "development":
            anaconda_dir = _update_conda_devel()
            print("Upgrading bcbio-nextgen to latest development version")
            pip_bin = os.path.join(os.path.dirname(sys.executable), "pip")
            git_tag = "@%s" % args.revision if args.revision != "master" else ""
            _pip_safe_ssl([[pip_bin, "install", "--upgrade", "--no-deps",
                            "git+%s%s#egg=bcbio-nextgen" % (REMOTES["gitrepo"], git_tag)]], anaconda_dir)
            print("Upgrade of bcbio-nextgen development code complete.")
        else:
            _update_conda_packages()
            print("Upgrade of bcbio-nextgen code complete.")

    try:
        _set_matplotlib_default_backend()
    except OSError:
        pass

    if args.tooldir:
        with bcbio_tmpdir():
            print("Upgrading third party tools to latest versions")
            _symlink_bcbio(args, script="bcbio_nextgen.py")
            _symlink_bcbio(args, script="bcbio_setup_genome.py")
            _symlink_bcbio(args, script="bcbio_prepare_samples.py")
            _symlink_bcbio(args, script="bcbio_fastq_umi_prep.py")
            upgrade_thirdparty_tools(args, REMOTES)
            print("Third party tools upgrade complete.")
    if args.toolplus:
        print("Installing additional tools")
        _install_toolplus(args)
    if args.install_data:
        for default in DEFAULT_INDEXES:
            if default not in args.aligners:
                args.aligners.append(default)
        if len(args.aligners) == 0:
            print("Warning: no aligners provided with `--aligners` flag")
        if len(args.genomes) == 0:
            print("Data not installed, no genomes provided with `--genomes` flag")
        else:
            with bcbio_tmpdir():
                print("Upgrading bcbio-nextgen data files")
                upgrade_bcbio_data(args, REMOTES)
                print("bcbio-nextgen data upgrade complete.")
    if args.isolate and args.tooldir:
        print("Isolated tool installation not automatically added to environmental variables")
        print(" Add:\n  {t}/bin to PATH".format(t=args.tooldir))
    save_install_defaults(args)
    args.datadir = _get_data_dir()
    _install_container_bcbio_system(args.datadir)
    print("Upgrade completed successfully.")
    return args

def _pip_safe_ssl(cmds, anaconda_dir):
    """Run pip, retrying with conda SSL certificate if global certificate fails.
    """
    try:
        for cmd in cmds:
            subprocess.check_call(cmd)
    except subprocess.CalledProcessError:
        _set_pip_ssl(anaconda_dir)
        for cmd in cmds:
            subprocess.check_call(cmd)

def _set_pip_ssl(anaconda_dir):
    """Set PIP SSL certificate to installed conda certificate to avoid SSL errors
    """
    if anaconda_dir:
        cert_file = os.path.join(anaconda_dir, "ssl", "cert.pem")
        if os.path.exists(cert_file):
            os.environ["PIP_CERT"] = cert_file

def _set_matplotlib_default_backend():
    """
    matplotlib will try to print to a display if it is available, but don't want
    to run it in interactive mode. we tried setting the backend to 'Agg'' before
    importing, but it was still resulting in issues. we replace the existing
    backend with 'agg' in the default matplotlibrc. This is a hack until we can
    find a better solution
    """
    if _matplotlib_installed():
        import matplotlib
        matplotlib.use('Agg', force=True)
        config = matplotlib.matplotlib_fname()
        with file_transaction(config) as tx_out_file:
            with open(config) as in_file, open(tx_out_file, "w") as out_file:
                for line in in_file:
                    if line.split(":")[0].strip() == "backend":
                        out_file.write("backend: agg\n")
                    else:
                        out_file.write(line)

def _matplotlib_installed():
    try:
        import matplotlib
    except ImportError:
        return False
    return True

def _symlink_bcbio(args, script="bcbio_nextgen.py"):
    """Ensure a bcbio-nextgen script symlink in final tool directory.
    """
    bcbio_anaconda = os.path.join(os.path.dirname(sys.executable), script)
    bindir = os.path.join(args.tooldir, "bin")
    if not os.path.exists(bindir):
        os.makedirs(bindir)
    bcbio_final = os.path.join(bindir, script)
    if not os.path.exists(bcbio_final):
        if os.path.lexists(bcbio_final):
            subprocess.check_call(["rm", "-f", bcbio_final])
        subprocess.check_call(["ln", "-s", bcbio_anaconda, bcbio_final])

def _install_container_bcbio_system(datadir):
    """Install limited bcbio_system.yaml file for setting core and memory usage.

    Adds any non-specific programs to the exposed bcbio_system.yaml file, only
    when upgrade happening inside a docker container.
    """
    base_file = os.path.join(datadir, "config", "bcbio_system.yaml")
    if not os.path.exists(base_file):
        return
    expose_file = os.path.join(datadir, "galaxy", "bcbio_system.yaml")
    expose = set(["memory", "cores", "jvm_opts"])
    with open(base_file) as in_handle:
        config = yaml.load(in_handle)
    if os.path.exists(expose_file):
        with open(expose_file) as in_handle:
            expose_config = yaml.load(in_handle)
    else:
        expose_config = {"resources": {}}
    for pname, vals in config["resources"].items():
        expose_vals = {}
        for k, v in vals.items():
            if k in expose:
                expose_vals[k] = v
        if len(expose_vals) > 0 and pname not in expose_config["resources"]:
            expose_config["resources"][pname] = expose_vals
    if expose_file and os.path.exists(os.path.dirname(expose_file)):
        with open(expose_file, "w") as out_handle:
            yaml.safe_dump(expose_config, out_handle, default_flow_style=False, allow_unicode=False)
    return expose_file

def _get_conda_bin():
    conda_bin = os.path.join(os.path.dirname(os.path.realpath(sys.executable)), "conda")
    if os.path.exists(conda_bin):
        return conda_bin

def _default_deploy_args(args):
    """Standard install arguments for CloudBioLinux.

    Avoid using sudo and keep an installation isolated if running as the root user.
    """
    return {"flavor": "ngs_pipeline_minimal",
            "vm_provider": "novm",
            "hostname": "localhost",
            "fabricrc_overrides": {"edition": "minimal",
                                   "use_sudo": False,
                                   "keep_isolated": args.isolate or os.geteuid() == 0,
                                   "conda_cmd": _get_conda_bin(),
                                   "distribution": args.distribution or "__auto__",
                                   "dist_name": "__auto__"}}

def _update_conda_packages():
    """If installed in an anaconda directory, upgrade conda packages.
    """
    conda_bin = _get_conda_bin()
    assert conda_bin, ("Could not find anaconda distribution for upgrading bcbio.\n"
                       "Using python at %s but could not find conda." % (os.path.realpath(sys.executable)))
    req_file = "bcbio-update-requirements.txt"
    if os.path.exists(req_file):
        os.remove(req_file)
    subprocess.check_call([conda_bin, "install", "--yes", "nomkl"])
    subprocess.check_call(["wget", "-O", req_file, "--no-check-certificate", REMOTES["requirements"]])
    subprocess.check_call([conda_bin, "install", "--update-deps", "--quiet", "--yes",
                           "-c", "bioconda", "-c", "conda-forge", "--file", req_file])
    if os.path.exists(req_file):
        os.remove(req_file)
    return os.path.dirname(os.path.dirname(conda_bin))

def _update_conda_devel():
    """Update to the latest development conda package.
    """
    conda_bin = _get_conda_bin()
    assert conda_bin, "Could not find anaconda distribution for upgrading bcbio"
    subprocess.check_call([conda_bin, "install", "--yes", "nomkl"])
    subprocess.check_call([conda_bin, "install", "--update-deps",
                           "--quiet", "--yes", "-c", "bioconda", "-c", "conda-forge", "bcbio-nextgen"])
    return os.path.dirname(os.path.dirname(conda_bin))

def get_genome_dir(gid, galaxy_dir, data):
    """Return standard location of genome directories.
    """
    if galaxy_dir:
        refs = genome.get_refs(gid, None, galaxy_dir, data)
        seq_file = tz.get_in(["fasta", "base"], refs)
        if seq_file and os.path.exists(seq_file):
            return os.path.dirname(os.path.dirname(seq_file))
    else:
        gdirs = glob.glob(os.path.join(_get_data_dir(), "genomes", "*", gid))
        if len(gdirs) == 1 and os.path.exists(gdirs[0]):
            return gdirs[0]

def _get_data_dir():
    base_dir = os.path.realpath(os.path.dirname(os.path.dirname(os.path.realpath(sys.executable))))
    if "anaconda" not in os.path.basename(base_dir) and "virtualenv" not in os.path.basename(base_dir):
        raise ValueError("Cannot update data for bcbio-nextgen not installed by installer.\n"
                         "bcbio-nextgen needs to be installed inside an anaconda environment \n"
                         "located in the same directory as `galaxy` `genomes` and `gemini_data` directories.")
    return os.path.dirname(base_dir)

def get_gemini_dir(data=None):
    try:
        data_dir = _get_data_dir()
        return os.path.join(data_dir, "gemini_data")
    except ValueError:
        if data:
            galaxy_dir = dd.get_galaxy_dir(data)
            data_dir = os.path.realpath(os.path.dirname(os.path.dirname(galaxy_dir)))
            return os.path.join(data_dir, "gemini_data")
        else:
            return None

def upgrade_bcbio_data(args, remotes):
    """Upgrade required genome data files in place.
    """
    from fabric.api import env
    data_dir = _get_data_dir()
    s = _default_deploy_args(args)
    s["actions"] = ["setup_biodata"]
    tooldir = args.tooldir or get_defaults().get("tooldir")
    if tooldir:
        s["fabricrc_overrides"]["system_install"] = tooldir
    s["fabricrc_overrides"]["data_files"] = data_dir
    s["fabricrc_overrides"]["galaxy_home"] = os.path.join(data_dir, "galaxy")
    cbl = get_cloudbiolinux(remotes)
    s["genomes"] = _get_biodata(cbl["biodata"], args)
    sys.path.insert(0, cbl["dir"])
    env.cores = args.cores
    cbl_deploy = __import__("cloudbio.deploy", fromlist=["deploy"])
    cbl_deploy.deploy(s)
    _upgrade_genome_resources(s["fabricrc_overrides"]["galaxy_home"],
                              remotes["genome_resources"])
    _upgrade_snpeff_data(s["fabricrc_overrides"]["galaxy_home"], args, remotes)
    if "vep" in args.datatarget:
        _upgrade_vep_data(s["fabricrc_overrides"]["galaxy_home"], tooldir)