#!/usr/bin/env python
# vim: set fileencoding=utf-8 :
#
# Copyright (C) 2010-2012 Marnix H. Medema
# University of Groningen
# Department of Microbial Physiology / Groningen Bioinformatics Centre
#
# Copyright (C) 2011,2012 Kai Blin
# University of Tuebingen
# Interfaculty Institute of Microbiology and Infection Medicine
# Div. of Microbiology/Biotechnology
#
# License: GNU Affero General Public License v3 or later
# A copy of GNU AGPL v3 should have been included in this software package in LICENSE.txt.

"""Script to download and install Pfam and ClusterBlast databases"""

import urllib2
import tarfile
import gzip
import backports.lzma as lzma
import hashlib
import os
import sys
import subprocess
import platform
import ctypes
from os import path

PFAM_URL = "ftp://ftp.ebi.ac.uk/pub/databases/Pfam/releases/Pfam27.0/Pfam-A.hmm.gz"
PFAM_CHECKSUM = "b29bc2c54db8090531df0361a781b8d7397f60ebedc0c36a16e7d45e999cc329"
CLUSTERBLAST_URL = "https://dl.secondarymetabolites.org/releases/4.0.0/clusterblast_20170105_v8_31.tar.xz"
CLUSTERBLAST_CHECKSUM = "c9c1f2c07ce97ea453345034727555928e7b0b8907f700805bb6a912865bb315"

LOCAL_FILE_PATH = os.path.abspath(os.path.dirname(__file__))

# pylint: disable=redefined-builtin
def execute(commands, input=None):
    "Execute commands in a system-independent manner"

    if input is not None:
        stdin_redir = subprocess.PIPE
    else:
        stdin_redir = None

    proc = subprocess.Popen(commands, stdin=stdin_redir,
                            stdout=subprocess.PIPE,
                            stderr=subprocess.PIPE)
    out, err = proc.communicate(input=input)
    retcode = proc.returncode
    return out, err, retcode
# pylint: enable=redefined-builtin


def get_remote_filesize(url):
    try:
        usock = urllib2.urlopen(url)
        dbfilesize = usock.info().get('Content-Length')
        if dbfilesize is None:
            dbfilesize = 0
    except urllib2.URLError:
        dbfilesize = 0
    dbfilesize = float(int(dbfilesize))  # db file size in bytes
    return dbfilesize


def get_free_space(folder):
    """ Return folder/drive free space (in bytes)"""
    if platform.system() == 'Windows':
        free_bytes = ctypes.c_ulonglong(0)
        ctypes.windll.kernel32.GetDiskFreeSpaceExW(ctypes.c_wchar_p(folder), None, None, ctypes.pointer(free_bytes))
        return free_bytes.value
    else:
        return os.statvfs(folder).f_bfree * os.statvfs(folder).f_frsize


def check_diskspace(file_url):
    """ Check if sufficient disk space is available"""
    dbfilesize = int(get_remote_filesize(file_url))
    free_space = int(get_free_space("."))
    if free_space < dbfilesize:
        print 'ERROR: Insufficient disk space available (required: %d, free: %d).' % (dbfilesize, free_space)
        sys.exit()


def download_file(url, filename):
    """ Download a file"""
    print "Downloading large file %s. Please be patient..." % (path.basename(filename))
    try:
        req = urllib2.urlopen(url)
    except urllib2.URLError:
        print 'ERROR: File not found on server.\nPlease check your internet connection.'
        sys.exit()
    CHUNK = 128 * 1024
    with open(filename, 'wb') as fp:
        while True:
            try:
                chunk = req.read(CHUNK)
                if not chunk:
                    break
                fp.write(chunk)
            except IOError:
                print 'ERROR: Download interrupted.'
                sys.exit()
    #Report download success
    print "Downloading %s finished successfully." % (path.basename(filename))
    return filename


def checksum(filename, chunksize=2 ** 20):
    """Get the SHA256 checksum of a file"""
    sha = hashlib.sha256()
    with open(filename, 'rb') as fh:
        for chunk in iter(lambda: fh.read(chunksize), b''):
            sha.update(chunk)

    return sha.hexdigest()


def unzip_file(filename, decompressor, error_type):
    """Decompress a compressed file"""
    newfilename, _ = path.splitext(filename)
    try:
        zipfile = decompressor.open(filename, 'rb')
        chunksize = 128 * 1024
        with open(newfilename, 'wb') as fp:
            while True:
                try:
                    chunk = zipfile.read(chunksize)
                    if not chunk:
                        break
                    fp.write(chunk)
                except IOError:
                    print 'ERROR: Unzipping interrupted.'
                    sys.exit()
    except error_type:
        print "ERROR: Error extracting %s. Please try to extract it manually." % (path.basename(filename))
        return
    print "Extraction of %s finished successfully." % (path.basename(filename))
    return newfilename


def untar_file(filename):
    """ Extract a TAR/GZ file"""
    try:
        tar = tarfile.open(filename)
        tar.extractall(path=filename.rpartition(os.sep)[0])
        tar.close()
    except tarfile.ReadError:
        print "ERROR: Error extracting %s. Please try to extract it manually." % (filename.rpartition(os.sep)[2])
        return
    print "Extraction of %s finished successfully." % (filename.rpartition(os.sep)[2])


def compile_pfam(filename):
    """ Compile a HMMer database with hmmpress"""
    command = ['hmmpress', '-f', filename]
    execute(command)


def delete_file(filename):
    try:
        os.remove(filename)
    except OSError:
        pass


def download_if_not_present(url, filename, sha256sum):
    """Download a file if it's not present or checksum doesn't match"""
    if path.exists(filename):
        print "Creating checksum of %s" % path.basename(filename)
        csum = checksum(filename)
        if csum == sha256sum:
            return
        else:
            print "checksum mismatch: expected %s, got %s" % (sha256sum, csum)

    download_file(url, filename)

    print "Creating checksum of %s" % path.basename(filename)
    csum = checksum(filename)
    if csum != sha256sum:
        print "Error downloading %s, sha256sum mismatch. Expected %s, got %s." % \
              (filename, sha256sum, csum)
        sys.exit(1)


def main():
    #Download and compile PFAM database
    check_diskspace(PFAM_URL)
    filename = path.join(LOCAL_FILE_PATH, "antismash", "generic_modules", "fullhmmer",
                         PFAM_URL.rpartition('/')[2])
    download_if_not_present(PFAM_URL, filename, PFAM_CHECKSUM)
    filename = unzip_file(filename, gzip, gzip.zlib.error)
    compile_pfam(filename)
    delete_file(filename + ".gz")

    #Download and compile ClusterBlast database
    check_diskspace(CLUSTERBLAST_URL)
    filename = path.join(LOCAL_FILE_PATH, "antismash", "generic_modules",
                         CLUSTERBLAST_URL.rpartition('/')[2])
    download_if_not_present(CLUSTERBLAST_URL, filename, CLUSTERBLAST_CHECKSUM)
    filename = unzip_file(filename, lzma, lzma.LZMAError)
    untar_file(filename)
    delete_file(filename)
    delete_file(filename + ".xz")

    # hmmpress the NRPS/PKS specific databases
    compile_pfam(path.join(LOCAL_FILE_PATH, "antismash", "specific_modules", "nrpspks", "abmotifs.hmm"))
    compile_pfam(path.join(LOCAL_FILE_PATH, "antismash", "specific_modules", "nrpspks", "dockingdomains.hmm"))
    compile_pfam(path.join(LOCAL_FILE_PATH, "antismash", "specific_modules", "nrpspks", "ksdomains.hmm"))
    compile_pfam(path.join(LOCAL_FILE_PATH, "antismash", "specific_modules", "nrpspks", "nrpspksdomains.hmm"))
    compile_pfam(path.join(LOCAL_FILE_PATH, "antismash", "specific_modules", "nrpspks", "sandpuma", "flat",
                           "fullset0_smiles", "fullset0_smiles_nrpsA.hmmdb"))

    # hmmpress the smcog specific database
    compile_pfam(path.join(LOCAL_FILE_PATH, "antismash", "generic_modules", "smcogs", "smcogs.hmm"))

    # hmmpress the ASF specific database
    compile_pfam(path.join(LOCAL_FILE_PATH, path.join("antismash", "generic_modules", "active_site_finder", "hmm", "p450.hmm3")))

    cmd = [path.join(LOCAL_FILE_PATH, 'scripts', 'create_bgc_seeds.sh'), LOCAL_FILE_PATH]
    _, err, retcode = execute(cmd)
    if retcode != 0:
        print "Error creating the BGC seeds profile"
        sys.exit(1)

    compile_pfam(path.join(LOCAL_FILE_PATH, "antismash", "generic_modules", "hmm_detection", "bgc_seeds.hmm"))

if __name__ == "__main__":
    main()
