#!/usr/bin/env python

# Longbow is Copyright (C) of James T Gebbie-Rayet and Gareth B Shannon 2015.
#
# This file is part of the Longbow software which was developed as part of
# the HECBioSim project (http://www.hecbiosim.ac.uk/).
#
# HECBioSim facilitates and supports high-end computing within the
# UK biomolecular simulation community on resources such as ARCHER.
#
# Longbow is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 2 of the License, or
# (at your option) any later version.
#
# Longbow is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Longbow.  If not, see <http://www.gnu.org/licenses/>.

"""Contains the main entry point for the Longbow program, this version calls
the main for a console based session."""

import os
import sys
import logging
import subprocess

# Depending on how longbow is installed/utilised the import will be slightly
# different, this should handle both cases.
try:

    APPLICATIONS = __import__("corelibs.applications", fromlist=[''])
    APPS = __import__("plugins.apps", fromlist=[''])
    CONFIGURATION = __import__("corelibs.configuration", fromlist=[''])
    EX = __import__("corelibs.exceptions", fromlist=[''])
    LOG = __import__("corelibs.logger", fromlist=[''])
    SCHEDULING = __import__("corelibs.scheduling", fromlist=[''])
    SHELLWRAPPERS = __import__("corelibs.shellwrappers", fromlist=[''])
    STAGING = __import__("corelibs.staging", fromlist=[''])

except ImportError:

    APPLICATIONS = __import__("Longbow.corelibs.applications", fromlist=[''])
    APPS = __import__("Longbow.plugins.apps", fromlist=[''])
    CONFIGURATION = __import__("Longbow.corelibs.configuration", fromlist=[''])
    EX = __import__("Longbow.corelibs.exceptions", fromlist=[''])
    LOG = __import__("Longbow.corelibs.logger", fromlist=[''])
    SCHEDULING = __import__("Longbow.corelibs.scheduling", fromlist=[''])
    SHELLWRAPPERS = __import__("Longbow.corelibs.shellwrappers", fromlist=[''])
    STAGING = __import__("Longbow.corelibs.staging", fromlist=[''])


def main(cmdline, files, mode, parameters):

    """This is the main for a console based app, this is designed to run in a
    python/unix shell and is thus the main choice for headless machines like a
    local cluster or for users that don't require a GUI."""

    # -------------------------------------------------------------------------
    # Determine some basic information

    # Get current directory (where we are run from).
    cwd = os.getcwd()

    # Get the execution directory (where we are installed).
    execdir = os.path.dirname(os.path.realpath(__file__))

    # -------------------------------------------------------------------------
    # Setup some basic file paths.

    try:

        # log
        # if a filename hasn't been provided default to log
        if files["log"] is "":

            files["log"] = "log"

        # if the path hasn't been provided default to the current working
        # directory
        if os.path.isabs(files["log"]) is False:

            files["log"] = os.path.join(cwd, files["log"])

        # ---------------------------------------------------------------------
        # Setup the logger.

        LOG.setuplogger(files["log"], "Longbow", mode)

        logger = logging.getLogger("Longbow")

        # ---------------------------------------------------------------------

        # hosts
        # if a filename hasn't been provided default to hosts.conf
        if files["hosts"] is "":

            files["hosts"] = "hosts.conf"

        # if the path hasn't been provided look in the current working
        # directory and then the execution directory if needs be
        if os.path.isabs(files["hosts"]) is False:

            if os.path.isfile(os.path.join(cwd, files["hosts"])):

                files["hosts"] = os.path.join(cwd, files["hosts"])

            elif os.path.isfile(os.path.join(execdir, files["hosts"])):

                files["hosts"] = os.path.join(execdir, files["hosts"])

            elif os.path.isfile(os.path.join(
                    os.path.expanduser("~/.Longbow"), files["hosts"])):

                files["hosts"] = os.path.join(
                    os.path.expanduser("~/.Longbow"), files["hosts"])

            else:

                raise EX.RequiredinputError(
                    "No host configuration file found in the current working "
                    "directory %s, the execution directory %s or in the "
                    "~/.Longbow directory." % (cwd, execdir))

        # job
        # if a job configuration file has been supplied but the path hasn't
        # look in the current working directory and then the execution
        # directory if needs be
        if files["job"] is not "" and os.path.isabs(files["job"]) is False:

            if os.path.isfile(os.path.join(cwd, files["job"])):

                files["job"] = os.path.join(cwd, files["job"])

            elif os.path.isfile(os.path.join(execdir, files["job"])):

                files["job"] = os.path.join(execdir, files["job"])

            else:

                raise EX.RequiredinputError(
                    "The job configuration file %s couldn't be found in the "
                    "current working directory %s, the execution directory %s."
                    % (files["job"], cwd, execdir))

    # -------------------------------------------------------------------------
    # setup and run some tests.

        # Log that we are starting up.
        logger.info("Welcome to Longbow!")
        logger.info("This software was developed as part of the EPSRC-funded"
                    " HECBioSim project (http://www.hecbiosim.ac.uk/)")
        logger.info("HECBioSim facilitates high-end biomolecular simulation "
                    "on resources such as ARCHER")
        logger.info("Longbow is Copyright (C) of James T Gebbie-Rayet and "
                    "Gareth B Shannon 2015.")
        logger.info("Longbow Commandline: {0}".format(cmdline))
        logger.info("Initialisation complete.")

        logger.info("hosts file is: %s", files["hosts"])

        # Load configurations and initialise Longbow data structures.
        hosts, jobs = CONFIGURATION.processconfigs(
            files["hosts"], files["job"], cwd, parameters)

        # Test the connection/s specified in the job configurations
        SHELLWRAPPERS.testconnections(hosts, jobs)

        # Test the hosts listed in the jobs configuration file have their
        # scheduler environments listed, if not then test and save them.
        SCHEDULING.testenv(files["hosts"], hosts, jobs)

        # Test that for the applications listed in the job configuration
        # file are available and that the executable is present.
        APPLICATIONS.testapp(hosts, jobs)

        # ---------------------------------------------------------------------
        # Start processing the setup and staging for each job.

        # Process the jobs command line arguments and find files for
        # staging.
        APPLICATIONS.processjobs(jobs)

        # Create jobfile and add it to the list of files that needs
        # uploading.
        SCHEDULING.prepare(hosts, jobs)

        # Stage all of the job files along with the scheduling script.
        STAGING.stage_upstream(hosts, jobs)

        # ---------------------------------------------------------------------
        # Submit the job/s to the scheduler.
        SCHEDULING.submit(hosts, jobs)

        # ---------------------------------------------------------------------
        # Monitor job/s.
        SCHEDULING.monitor(hosts, jobs)

    # -------------------------------------------------------------------------
    # Handle the errors and Longbow exit.

    except (EX.RsyncError, EX.SCPError, EX.SSHError) as err:

        # Output the information about the problem.
        if mode["debug"]:

            logger.exception(err)

        else:

            logger.error(err)

        logger.error("stdout:\n\n{}" .format(err.stdout))
        logger.error("stderr:\n\n{}" .format(err.stderr))
        logger.error("errorcode: {}" .format(err.errorcode))

    except (SystemExit, KeyboardInterrupt):

        logger.info("User interrupt - performing cleanup.")

        # User has decided to kill Longbow, check if there are any jobs still
        # running.
        for job in jobs:

            if "jobid" in jobs[job]:

                # If job is not finished delete and stage.
                if (jobs[job]["laststatus"] != "Finished" and
                        jobs[job]["laststatus"] != "Submit Error"):

                    # Kill it.
                    SCHEDULING.delete(hosts, jobs, job)

                    # Transfer the directories as they are.
                    STAGING.stage_downstream(hosts, jobs, job)

                # Job is finished then just stage.
                elif jobs[job]["laststatus"] != "Submit Error":

                    # Transfer the directories as they are.
                    STAGING.stage_downstream(hosts, jobs, job)

    except EX.RequiredinputError as err:

        logger.error(err)

    except Exception as err:

        if mode["debug"]:

            logger.exception(err)

        else:

            logger.error(err)

    finally:

        # Cleanup.

        # There is a case where when longbow is launched without any args that
        # a glitch here will suppress all the error messages from being seen.
        try:

            STAGING.cleanup(hosts, jobs)

        except NameError:

            pass

        logger.info("Good bye from Longbow!")
        logger.info("Check out http://www.hecbiosim.ac.uk/ for other " +
                    "powerful biomolecular simulation software tools.")

    # -------------------------------------------------------------------------


if __name__ == "__main__":

    """Main entry point for Longbow.

    To run Longbow, simply write longbow before the command you wish to
    be executed using your chosen simulation package e.g.:

    %longbow pmemd.MPI -i example.in -c example.min -p example.top -o output

    In addition, the following longbow flags may be provided before the
    executable (pmemd.MPI in the above example):

    -hosts filename
    -job filename
    -log filename
    -resource remote resource name
    -replicates number
    -jobname jobname
    -debug
    -verbose
    -about
    -version
    -help
    -examples

    Read the documentation at http://www.hecbiosim.ac.uk/ for more information
    on how to setup and run jobs using Longbow.
    """

    # ------------------------------------------------------------------------
    # Some defaults.

    # Actual command line.
    COMMANDLINE = (" ").join(sys.argv)

    # Fetch command line arguments as list and remove longbow exec
    COMMANDLINEARGS = sys.argv
    COMMANDLINEARGS.pop(0)

    # Initialise file path params, so we can pass blank to signify use default
    # paths if not supplied.
    FILES = {
        "hosts": "",
        "job": "",
        "log": ""
    }

    # Initialise parameters that could alternatively be provided in
    # configuration files
    PARAMETERS = {
        "resource": "",
        "replicates": "",
        "jobname": "",
        "executable": "",
        "executableargs": ""
    }

    # Initialise output modes
    MODE = {
        "debug": False,
        "verbose": False
    }

    # Specify all recognised longbow arguments
    ALLLONGBOWARGS = [
        "-hosts",
        "-job",
        "-log",
        "-resource",
        "-replicates",
        "-jobname",
        "-debug",
        "-verbose",
        "-about",
        "-version",
        "--version",
        "-V",
        "-help",
        "--help",
        "-h",
        "-examples"
    ]

    # Initialise
    EXECUTABLE = ""
    POSITION = ""
    LONGBOWARGS = ""
    EXECARGS = ""

    # ------------------------------------------------------------------------
    # Split the command line into longbow arguments, the executable and
    # executable arguments

    # Get a list of recognised executables
    execlist = getattr(APPS, "EXECLIST")

    # Search for recognised executables on the commandline
    for exe in execlist:

        if exe in COMMANDLINEARGS:

            if COMMANDLINEARGS.count(exe) == 1:

                EXECUTABLE = exe
                POSITION = COMMANDLINEARGS.index(exe)
                LONGBOWARGS = COMMANDLINEARGS[:POSITION]
                EXECARGS = COMMANDLINEARGS[POSITION+1:]
                break

            else:

                raise EX.CommandlineargsError(
                    "More than one recognised executable has been specified on"
                    " the Longbow command line. Please provide just one.")

    # if an executable wasn't found perhaps it is specified in a configuration
    # file. Executable arguments might still be provided on the command line so
    # try to detect them.
    if EXECUTABLE == "":

        for item in COMMANDLINEARGS:

            POSITION = COMMANDLINEARGS.index(item)

            # if item provided on the commandline doesn't appear to be a
            # longbow argument, assume it, and all others following it are
            # executable arguments
            if (item not in ALLLONGBOWARGS and
                    COMMANDLINEARGS[POSITION-1][1:] not in FILES and
                    COMMANDLINEARGS[POSITION-1][1:] not in PARAMETERS):

                LONGBOWARGS = COMMANDLINEARGS[:POSITION]
                EXECARGS = COMMANDLINEARGS[POSITION:]
                break

        # if no executable arguments have been found, everything on the command
        # line must be longbow arguments since we know the executable hasn't
        # been provided on the command line either.
        if EXECARGS == "":

            LONGBOWARGS = COMMANDLINEARGS

    # Make sure the user hasn't provided bogus longbow arguments or those
    # that aren't recognised on the command line
    for item in LONGBOWARGS:

        if item.startswith("-") and item not in ALLLONGBOWARGS:

            allowedargs = " ".join(ALLLONGBOWARGS)

            raise EX.CommandlineargsError(
                "Argument %s is not a recognised Longbow argument. "
                "Recognised arguments are: %s" % (item, allowedargs))

    # ------------------------------------------------------------------------
    # Determine if the user is trying to run a sub-functionality such as
    # % longbow -about for example.

    # Take test for the about command line flag, print message and exit if
    # found
    if LONGBOWARGS.count("-about") == 1:

        print (
            "Welcome to Longbow!\n"
            "Longbow is a remote job submission utility designed for "
            "biomolecular simulation."
            "This software was developed as part of the EPSRC-funded "
            "HECBioSim project \n" "(http://www.hecbiosim.ac.uk/).\n"
            "HECBioSim facilitates high-end biomolecular simulation "
            "on resources such as \n" "ARCHER. \n"
            "Longbow is Copyright (C) of James T Gebbie-Rayet and "
            "Gareth B Shannon 2015.")

        exit(0)

    # Take test for the version command line flag, print message and exit
    # if found
    elif (LONGBOWARGS.count("-version") == 1 or
          LONGBOWARGS.count("--version") == 1 or LONGBOWARGS.count("-V") == 1):

        print ("Longbow 1.01.004")

        exit(0)

    # Take test for the help command line flag, print message and exit if
    # found
    elif (LONGBOWARGS.count("-help") == 1 or
          LONGBOWARGS.count("--help") == 1 or
          LONGBOWARGS.count("-h") == 1):

        print (
            "Usage:\n\n"
            "Before running Longbow, first setup a passwordless\n"
            "connection with a target remote resource and setup\n"
            "configuration files according to the documentation. Submit jobs\n"
            "using the following format:\n\n"
            "longbow [longbow args] executable [executable args]\n\n"
            "e.g.:\n"
            "%longbow -verbose pmemd.MPI -i example.in -c example.min -p "
            "example.top -o output\n\n"
            "longbow args:\n\n"
            "-hosts filename           : specifies the hosts configuration "
            "file filename\n"
            "-job filename             : specifies the job configuration file"
            " filename\n"
            "-log filename             : specifies the file Longbow output"
            " should be directed to\n"
            "-resource remote resource : specifies the remote resource\n"
            "-replicates number        : number of replicate jobs to be "
            "submitted\n"
            "-jobname jobname          : the name of the job to be submitted\n"
            "-debug                    : additional output to assist debugging"
            "\n"
            "-verbose                  : additional run-time info to be output"
            "\n"
            "-about                    : prints Longbow description\n"
            "-version, --version, -V   : prints Longbow version number\n"
            "-help, --help, -h         : prints Longbow help\n"
            "-examples                 : downloads example files to "
            "./LongbowExamples"
            "\n\n"
            "Read the documentation at http://www.hecbiosim.ac.uk/ for \n"
            "more information on how to setup and run jobs using Longbow.")

        exit(0)

    # Take test for the examples command line flag, download files and exit
    # if found
    elif LONGBOWARGS.count("-examples") == 1:

        if not os.path.isfile(
                os.path.join(os.getcwd(), "LongbowExamples.zip")):

            try:

                subprocess.check_output([
                    "wget",
                    "http://www.hecbiosim.ac.uk/longbow/send/5-longbow/" +
                    "4-longbow-examples",
                    "-O",
                    os.path.join(os.getcwd(), "LongbowExamples.zip")])

            except subprocess.CalledProcessError:

                subprocess.call([
                    "curl",
                    "-L",
                    "http://www.hecbiosim.ac.uk/longbow/send/5-longbow/" +
                    "4-longbow-examples",
                    "-o",
                    os.path.join(os.getcwd(), "LongbowExamples.zip")])

            subprocess.call(["unzip", "-d", os.getcwd(),
                             os.path.join(os.getcwd(), "LongbowExamples.zip")])

        exit(0)

    # ------------------------------------------------------------------------
    # Pull out some of the specific commandline args

    # Store the config file path
    if LONGBOWARGS.count("-hosts") == 1:

        POSITION = LONGBOWARGS.index("-hosts")

        if (POSITION + 1 == len(LONGBOWARGS) or
                LONGBOWARGS[POSITION + 1].startswith("-")):

            raise EX.CommandlineargsError(
                "Please specify a valid file for the -hosts command line"
                " parameter e.g. longbow -hosts filename ...")
        else:

            FILES["hosts"] = LONGBOWARGS[POSITION + 1]

    # Store the job config file path
    if LONGBOWARGS.count("-job") == 1:

        POSITION = LONGBOWARGS.index("-job")

        if (POSITION + 1 == len(LONGBOWARGS) or
                LONGBOWARGS[POSITION + 1].startswith("-")):

            raise EX.CommandlineargsError(
                "Please specify a valid file for the -job command line"
                " parameter e.g. longbow -job filename ...")

        else:

            FILES["job"] = LONGBOWARGS[POSITION + 1]

    # Store the log file path
    if LONGBOWARGS.count("-log") == 1:

        POSITION = LONGBOWARGS.index("-log")

        if (POSITION + 1 == len(LONGBOWARGS) or
                LONGBOWARGS[POSITION + 1].startswith("-")):

            raise EX.CommandlineargsError(
                "Please specify a valid file for the -log command line"
                " parameter e.g. longbow -log filename ...")

        else:

            FILES["log"] = LONGBOWARGS[POSITION + 1]

    # Store the DEBUG parameter
    if LONGBOWARGS.count("-debug") == 1:

        POSITION = LONGBOWARGS.index("-debug")
        MODE["debug"] = True

    # Store the VERBOSE parameter
    if LONGBOWARGS.count("-verbose") == 1:

        POSITION = LONGBOWARGS.index("-verbose")
        MODE["verbose"] = True

    # Store the resource name
    if LONGBOWARGS.count("-resource") == 1:

        POSITION = LONGBOWARGS.index("-resource")

        if (POSITION + 1 == len(LONGBOWARGS) or
                LONGBOWARGS[POSITION + 1].startswith("-")):

            raise EX.CommandlineargsError(
                "Please specify a valid remote resource for the -resource"
                " command line parameter e.g. longbow -resource HPC ...")

        else:

            PARAMETERS["resource"] = LONGBOWARGS[POSITION + 1]

    # Store the replicates name
    if LONGBOWARGS.count("-replicates") == 1:

        POSITION = LONGBOWARGS.index("-replicates")

        if (POSITION + 1 == len(LONGBOWARGS) or
                LONGBOWARGS[POSITION + 1].startswith("-")):

            raise EX.CommandlineargsError(
                "Please specify a number for the -replicates command line"
                " parameter e.g. longbow -replicates 5 ...")

        else:

            PARAMETERS["replicates"] = LONGBOWARGS[POSITION + 1]

    # Store the jobname name
    if LONGBOWARGS.count("-jobname") == 1:

        POSITION = LONGBOWARGS.index("-jobname")

        if (POSITION + 1 == len(LONGBOWARGS) or
                LONGBOWARGS[POSITION + 1].startswith("-")):

            raise EX.CommandlineargsError(
                "Please specify a name for the -jobname command line parameter"
                " e.g. longbow -jobname myjob ...")

        else:

            PARAMETERS["jobname"] = LONGBOWARGS[POSITION + 1]

    PARAMETERS["executable"] = EXECUTABLE
    PARAMETERS["executableargs"] = EXECARGS

    # ------------------------------------------------------------------------
    # Call Longbow.

    # Enter the mains application.
    main(COMMANDLINE, FILES, MODE, PARAMETERS)
