#!/usr/bin/env python3
#------------------------------------------------------------------------------#
#  DFTB+: general package for performing fast atomistic simulations            #
#  Copyright (C) 2006 - 2023  DFTB+ developers group                           #
#                                                                              #
#  See the LICENSE file for terms of usage and distribution.                   #
#------------------------------------------------------------------------------#
#
'''
Downloads optional external components
'''
from __future__ import print_function
import sys
import os
import argparse
import subprocess as sub
import shutil

USAGE = '''%(prog)s [options] [OPTEXT [OPTEXT ...]]

Downloads optional external components.

Possible choices for OPTEXT:

  UNRESTRICTED (default) -- download only components which are less or equally
      restrictive than the LGPLv3 license of DFTB+ and do not restrict the
      distribution of the compiled code beyond its LGPL-license.

  RESTRICTED -- download optional components with stricter license than the
      LGPLv3 license of DFTB+ (typically GPL). The distribution of the compiled
      code will be restricted beyond the LGPL license and governed by the
      strictest license of the included optional components (typically GPLv3).

  ALL -- download all optional external components.

      NOTE: this will also download components which may restrict the license of
      the combined code beyond the LGPLv3 license of DFTB+.

  slakos -- Parametrisations for testing DFTB+ [license: CC-BY-SA, UNRESTRICTED]

  gbsa -- Parametrisations for the GBSA model [license: CC-BY-SA, UNRESTRICTED]
'''

RESTRICTED_WARNING = """WARNING:
    This optional external component has a more restrictive license than the
    DFTB+ license. Restrictions on distribution of the compiled code may apply.
"""

SCRIPT_DIR = os.path.dirname(sys.argv[0])
EXTERNAL_DIR = os.path.abspath(os.path.join(SCRIPT_DIR, '..', 'external'))
STATUS_FILE = 'status'

STATUS_MISSING = 'missing'
STATUS_OK = 'ok'

UNRESTRICTED_EXTERNAL = 0
RESTRICTED_EXTERNAL = 1


class ScriptError(Exception):
    '''Exception raised by the script.'''
    pass


def main():
    '''Main routine.'''

    opts, args = parse_cmdline_args()
    excluded = set(opts.excluded) if opts.excluded is not None else set()

    externals = []
    for external in args:
        if external == 'UNRESTRICTED':
            externals += EXTERNALS_UNRESTRICTED
        elif external == 'RESTRICTED':
            externals += EXTERNALS_RESTRICTED
        elif external == 'ALL':
            externals += EXTERNALS
        else:
            externals.append(external)

    done = set()
    for external in externals:
        if external in done or external in excluded:
            continue
        getter, destdir, restriction = GETTERS[external]
        destdir = os.path.join(EXTERNAL_DIR, destdir)
        if not os.path.isdir(destdir):
            msg = "Directory '{}' does not exist".format(destdir)
            raise ScriptError(msg)
        print('\n*** Getting external:', external, '\n    into', destdir, '\n')
        if restriction == RESTRICTED_EXTERNAL:
            print(RESTRICTED_WARNING)
        set_status(destdir, STATUS_MISSING)
        getter(destdir)
        set_status(destdir, STATUS_OK)
        print()
        done.add(external)


def parse_cmdline_args():
    '''Parses and returns command line arguments.'''
    parser = argparse.ArgumentParser(usage=USAGE)
    msg = 'Exclude certain externals (mostly useful when downloading all '\
          'optional externals except a few)'
    parser.add_argument(
        '-e', '--exclude', dest='excluded', metavar='EXCLUDED',
        choices=EXTERNALS, default=None, action='append', help=msg)
    opts, args = parser.parse_known_args()

    argchoices = ['UNRESTRICTED', 'RESTRICTED', 'ALL'] + EXTERNALS
    if not args:
        args = ['UNRESTRICTED']
    for arg in args:
        if arg not in argchoices:
            msg = "Invalid optional external '{}'".format(arg)
            parser.error(msg)
    return opts, args


def set_status(destdir, status):
    '''Update status file.'''
    statusfile = os.path.join(destdir, STATUS_FILE)
    with open(statusfile, 'w') as fp:
        fp.write(status)


def get_slakos(destdir):
    '''Downloads and extracts parametrizations for the tests.'''

    url = 'https://github.com/dftbplus/testparams/archive/6165104e60efbdb3ad05d005c282ada50f12dfef.tar.gz'
    fname = os.path.join(destdir, 'testparams.tar.gz')
    unpackname = os.path.join(destdir, 'testparams-6165104e60efbdb3ad05d005c282ada50f12dfef')
    targetname = os.path.join(destdir, 'origin')
    if os.path.exists(unpackname):
        shutil.rmtree(unpackname)
    if os.path.exists(targetname):
        shutil.rmtree(targetname)
    download_file(url, fname)
    unpack_file(fname)
    os.remove(fname)
    os.rename(unpackname, targetname)


def get_gbsa(destdir):
    '''Downloads and extracts parametrizations for gbsa.'''

    url = 'https://github.com/grimme-lab/gbsa-parameters/archive/6836c4d997e4135e418cfbe273c96b1a3adb13e2.tar.gz'
    fname = os.path.join(destdir, 'gbsa-parameters.tar.gz')
    unpackname = os.path.join(destdir, 'gbsa-parameters-6836c4d997e4135e418cfbe273c96b1a3adb13e2')
    targetname = os.path.join(destdir, 'origin')
    if os.path.exists(unpackname):
        shutil.rmtree(unpackname)
    if os.path.exists(targetname):
        shutil.rmtree(targetname)
    download_file(url, fname)
    unpack_file(fname)
    os.remove(fname)
    os.rename(unpackname, targetname)


def download_file(url, fname):
    '''Download file (currently via wget)'''
    print('Download:\n    {0}\n    from {1}'.format(fname, url))
    sub.check_call(['wget', '-q', '-O', fname, url])


def unpack_file(fname):
    '''Unpack file (currently via tar)'''
    fdir = os.path.dirname(fname)
    if fname.endswith('.tar.xz'):
        sub.check_call(['tar', '-x', '-J', '-C', fdir, '-f', fname])
    elif fname.endswith('.tar.bz2'):
        sub.check_call(['tar', '-x', '-j', '-C', fdir, '-f', fname])
    elif fname.endswith('.tar.gz'):
        sub.check_call(['tar', '-x', '-z', '-C', fdir, '-f', fname])
    elif fname.endswith('.tar'):
        sub.check_call(['tar', '-x', '-C', fdir, '-f', fname])
    else:
        msg = "File '{}' has unknown archive suffix".format(fname)
        raise ScriptError(msg)


def query_yes_no(question, default=True):
    """Ask a yes/no question via raw_input() and return True/False.

    Args:
        question: string presented to the user.
        default: is the presumed answer if the user just hits <Enter>.
            It must be "yes" (the default), "no" or None (meaning
            an answer is required of the user).

    Returns:
        The "answer" return value is True for "yes" or False for "no".
    """
    valid = {"yes": True, "y": True, "no": False, "n": False}
    if default is None:
        prompt = " [y/n] "
    elif default:
        prompt = " [Y/n] "
        defaultstr = "y"
    else:
        prompt = " [y/N] "
        defaultstr = "n"

    while True:
        sys.stdout.write(question + prompt)
        choice = read_input().lower()
        if default is not None and choice == '':
            return valid[defaultstr]
        elif choice in valid:
            return valid[choice]
        else:
            sys.stdout.write("Please respond with 'yes' or 'no' "
                             "(or 'y' or 'n').\n")


def query_download_confirmation(component, default=False):
    '''Confirmation query for downloading a component.'''
    question = '    Download {0}?'.format(component)
    answer = query_yes_no(question, default)
    print()
    return answer


def read_input():
    '''Input method working with both Python 2 and 3.'''
    if sys.version_info[0] >= 3:
        return input()
    else:
        return raw_input()


# Methods to obtain externals.
# Contains (function, targetfolder, restriction) tuples
GETTERS = {
    'slakos': (get_slakos, 'slakos', UNRESTRICTED_EXTERNAL),
    'gbsa': (get_gbsa, 'gbsa', UNRESTRICTED_EXTERNAL),
}

# All externals
EXTERNALS = list(GETTERS.keys())

# LGPL-compatible externals without restriction on distribution
EXTERNALS_UNRESTRICTED = [ext for ext in GETTERS.keys()
                          if GETTERS[ext][2] == UNRESTRICTED_EXTERNAL]

# LGPL-compatible externals with restriction on distribution
EXTERNALS_RESTRICTED = [ext for ext in GETTERS.keys()
                        if GETTERS[ext][2] == RESTRICTED_EXTERNAL]


if __name__ == '__main__':
    main()
