#! /usr/bin/env python3

import click
import contextlib
import distutils.spawn
import os
import subprocess
import sys
import tempfile
import textwrap

import job_stream.usersettings as usersettings

ENVVAR = "JOBS_CONFIG"
ENVVAR_CHECKPOINT = "JOBS_CHECKPOINT"

def _checkpointArgs(c, checkpoint):
    chkpt = None
    if c:
        if checkpoint is not None:
            raise ValueError("Use one of -c or --checkpoint, not both")
        chkpt = "job_stream.chkpt"
    elif checkpoint is not None:
        chkpt = checkpoint

    return chkpt


def _settings():
    s = usersettings.Settings('com.lamegameproductions.job_stream')
    s.add_setting("hostfile", str, default=None)
    s.load_settings()

    # Stored as string
    if s.hostfile == 'None':
        s.hostfile = None

    return s


def _which(prog):
    return distutils.spawn.find_executable(prog)


def doRun(hostfile, checkpoint, argv):
    """Runs the program and arguments specified by argv using mpirun with
    hostfile.
    """
    mpirun = _which('mpirun')
    prog = _which(argv[0])
    useMpi = True
    env = os.environ.copy()
    env['PYTHONUNBUFFERED'] = '1'
    env[ENVVAR] = hostfile
    if checkpoint is not None:
        env[ENVVAR_CHECKPOINT] = checkpoint

        if os.path.lexists(checkpoint + '.done'):
            # job_stream WILL exit anyway.  Therefore, do not run it over
            # MPI and pollute stderr with a bunch of copies of the same
            # message
            useMpi = False

    hosts = []  # List, not set, so that order is preserved
    with open(hostfile) as f:
        for line in f:
            l = line.strip()
            if not l or l.startswith("#"):
                continue
            host, other = (l + ' ').split(' ', 1)
            if host in hosts:
                raise ValueError("Duplicate host entry? {}".format(host))
            hosts.append(host)
    hosts = ','.join(hosts)

    # MPI does not automatically export environment variables... ugh.
    # Supposedly this makes sense in certain super computing contexts, however
    # not exporting environment variables would make job_stream applications
    # need to follow a much more explicit format.  The idea is that you can
    # execute the program locally to test, or with `job_stream` for
    # parallelization.
    mpiEnvVars = [ ('-x', key) for key in env ]
    mpiArgs = (
            (mpirun, '-host', hosts)
            + ('--map-by', 'node', '--bind-to', 'none')
            + ('--mca', 'orte_abort_on_non_zero_status', '1')
            + tuple(b for a in mpiEnvVars for b in a))
    if not useMpi:
        mpiArgs = ()

    r = subprocess.Popen(mpiArgs + (prog,) + argv[1:], env=env)
    sys.exit(r.wait())


@click.command()
@click.option('-c/--no-c', help="Enable checkpoints with the default name "
        "job_stream.chkpt.  Note that this file will reside in the current "
        "working directory.")
@click.option('--checkpoint', help="Enable checkpoints with a custom "
        "name/path; may also be supplied as {}.".format(ENVVAR_CHECKPOINT),
        envvar=ENVVAR_CHECKPOINT)
@click.option('--hostfile', help="Use HOSTFILE to determine which machines "
        "should run the job_stream process.")
@click.option('--host', help="Run the job_stream process on HOST, a "
        "comma-separated list of servers, the same as lines from a hostfile.  "
        "E.g., --host 'cpu1,cpu2 maxCpu=2'.")
@click.argument('prog')
@click.argument('args', nargs=-1)
def main(prog, args, hostfile, host, c, checkpoint):
    """
    Runs the job_stream application PROG using ARGS as arguments.  Note that
    PROG will be resolved to an absolute path from within the current
    environment; for Python applications, this means that running job_stream
    within a virtualenv is safe.

    Often, PROG should be preceded by '--', to let the option parser know that
    any options after PROG belong to PROG and not to job_stream (e.g.
    job_stream -- bash -c 'echo `hostname`').

    PROG may be 'config' to access job_stream's default configuration.

    PROG may be 'git-results-progress' to print the timestamp of the checkpoint
    file associated with job_stream's arguments.  This will raise an exception
    if no checkpoint argument (-c, --checkpoint, or JOBS_CHECKPOINT) is
    specified.
    """
    settings = _settings()

    ## Hostfile configuration
    thisHostfile = None
    @contextlib.contextmanager
    def tmpHostfile():
        """Empty context block since normal usage has no tmpfile"""
        yield
    tmpHostfile = tmpHostfile()

    if hostfile is not None:
        thisHostfile = os.path.abspath(hostfile)
        if host is not None:
            raise ValueError("Incompatible options: --hostfile and --host")
    elif host is not None:
        tmpHostfile = tempfile.NamedTemporaryFile("w+t")
        tmpHostfile.write("\n".join(host.split(",")))
        tmpHostfile.flush()

        thisHostfile = tmpHostfile.name

    if thisHostfile is None:
        if settings.hostfile is None:
            sys.stderr.write("No --hostfile specified and no --default has "
                    "been set\n")
            try:
                main(['--help'])
            except SystemExit:
                pass
            sys.exit(2)

        thisHostfile = settings.hostfile

    ## Checkpoint configuration
    chkpt = _checkpointArgs(c, checkpoint)

    ## Checkpoint progress check?
    if prog.lower() == 'git-results-progress':
        gitResultsProgress(chkpt, args)
    else:
        with tmpHostfile:
            doRun(thisHostfile, chkpt, (prog,) + args)


@click.command()
@click.argument('setting', nargs=-1)
def config(setting):
    r"""With no SETTING, prints out all configuration options available.

    Otherwise, each SETTING is a KEY=VALUE to set for this user's
    configuration.

    \b
    Valid KEYs:
        hostfile: Path to the default hostfile to use (--hostfile)
    """
    keys = [ 'hostfile' ]

    settings = _settings()
    if not setting:
        for k in keys:
            print("{}={}".format(k, getattr(settings, k)))
    else:
        for s in setting:
            if '=' not in s:
                raise ValueError("Each SETTING must have an '=' in it to "
                        "distinguish between KEY and VALUE: {}".format(s))

            key, val = s.split('=', 1)
            if key not in keys:
                raise ValueError("KEY '{}' not found?  {}".format(key, keys))

            v = val.strip()
            if not v:
                v = None

            if key == 'hostfile':
                # Path
                setattr(settings, key,
                        os.path.abspath(v) if v is not None else v)
            else:
                # Other
                setattr(settings, key, v)
        settings.save_settings()


def gitResultsProgress(chkpt, args):
    """Utility for git-results' "progress" configuration for auto-retry, or
    for any other utility that wants to know if the job_stream is making
    progress or failing repeatedly.

    Prints, to stdout, a single, monotonically-increasing floating-point number
    that will be a constant when progress is not being made.
    """
    if args:
        print("Usage: git-results [-c] [--checkpoint CHECKPOINT] "
                "git-results-progress\n\n{}".format(
                    gitResultsProgress.__doc__))
        sys.exit(0)

    if chkpt is None:
        raise ValueError("No checkpoint file specified")

    try:
        print(os.path.getmtime(chkpt))
    except FileNotFoundError:
        try:
            print(os.path.getmtime(chkpt + ".done"))
        except FileNotFoundError:
            print("-1")


if __name__ == '__main__':
    if len(sys.argv) < 2:
        main(['--help'])
    elif sys.argv[1].lower() == 'config':
        config(sys.argv[2:])
    else:
        main(sys.argv[1:])

