#!/usr/bin/env python3
#
# Command-Line Interface for running tasks
#

import datetime
import optparse
import os
import shlex
import pscheduler
import subprocess
import sys
import time


pscheduler.set_graceful_exit()


#
# Gargle the arguments
#


class VerbatimParser(optparse.OptionParser):
    def format_epilog(self, formatter):
        return self.epilog

opt_parser = VerbatimParser(
    usage="Usage: %prog [options] test-type [test-options]",
    epilog=
"""
For help with a test:

  task test-type --help  (e.g., task trace --help)

Examples:

  task rtt --dest ps.example.com
      Round-trip test from here to ps.example.com

  task --repeat PT30M --max-runs 5 rtt --dest ps.example.com
      Same, repeated every 30 minutes up to five times

  task --tool tracepath trace --dest ps.example.com
      Trace test using the tool "tracepath"

  task --export rtt --count 10 --length 128 --dest ps.example.com > mytask
      Export JSON describing task to file "mytask"

  task --import mytest
      Import and run the task described in the file "mytask"

  task --import mytask - --dest ps.example.net
      Import the task described in the file "mytask", change the destination
      to ps.example.net and run it.  Note that the single-hyphen test type
      is required if any parameter changes are to be added on the command
      line.  Specifying a test type here will override the imported value.
"""
    )
opt_parser.disable_interspersed_args()


# TASK OPTIONS

task_group = optparse.OptionGroup(opt_parser, "Task Options")
opt_parser.add_option_group(task_group)

task_group.add_option("--archive",
                      help="Specify where to archive result(s) (JSON; optionally @/path/to/file; may be repeated)",
                      default=[],
                      action="append", type="string",
                      dest="archive")

task_group.add_option("--keep-after-archive",
                      help="Maximim time to keep runs after archiving is complete (ISO 8601 Duration)",
                      action="store",
                      type="string",
                      dest="keep_after_archive")

task_group.add_option("--context",
                      help="Specify context changes for all participants (JSON; optionally @/path/to/file)",
                      default=None,
                      action="store", type="string",
                      dest="context")

task_group.add_option("--key",
                      help="Key required for write access to the task (Optional @/path/to/file)",
                      action="store", type="string",
                      dest="key")

task_group.add_option("--reference",
                      help="Save arbitrary JSON with task for reference (Optional @/path/to/file)",
                      action="store", type="string",
                      dest="reference")

task_group.add_option("--tool",
                      help="Choose a tool to use for the test (May be repeated for a preferred-order selection)",
                      default=[],
                      action="append", type="string",
                      dest="tool")



# SCHEDULING OPTIONS

schedule_group = optparse.OptionGroup(opt_parser, "Scheduling Options")
opt_parser.add_option_group(schedule_group)

schedule_group.add_option("--max-runs",
                      help="Maximum number of repeats (requires --repeat)",
                      action="store", type="int", default=1,
                      dest="max_runs")

task_group.add_option("--priority",
                      help="Scheduling priority to request (default 0 or $PSCHEDULER_PRIORTY)",
                      action="store", type="int",
                      dest="priority")

schedule_group.add_option("--repeat",
                      help="Repeat interval (ISO 8601 Duration)",
                      action="store", type="string",
                      dest="repeat")

schedule_group.add_option("--repeat-cron",
                          help="Cron-style repeat specification (POSIX Cron)",
                          action="store", type="string",
                          dest="repeat_cron")

schedule_group.add_option("--slip",
                      help="Allowed start slip (ISO8601 Duration)",
                      action="store", type="string",
                      dest="slip")

schedule_group.add_option("--sliprand",
                      help="Slip randomly",
                      action="store_true",
                      dest="sliprand")

schedule_group.add_option("--start",
                      help="Start time",
                      action="store", type="string",
                      dest="start")

schedule_group.add_option("--until",
                      help="Time after which scheduling should stop",
                      action="store", type="string",
                      dest="until")


# OTHER OPTIONS

other_group = optparse.OptionGroup(opt_parser, "Other Options")
opt_parser.add_option_group(other_group)

other_group.add_option("--assist",
                      help="Use the host ASSIST for assistance (default localhost or $PSCHEDULER_ASSIST)",
                      action="store", type="string",
                      dest="assist")

other_group.add_option("--assist-bind",
                       help="Bind to the supplied address when looking for assistance" \
                       "(default localhost or $PSCHEDULER_ASSIST_BIND)",
                      action="store", type="string",
                      dest="assist_bind")

other_group.add_option("--bind",
                       help="Bind to the supplied address when submitting the task",
                       action="store", type="string",
                       default=None,
                       dest="bind")

other_group.add_option("--export",
                      help="Export (unvalidated) task JSON to stdout and exit",
                      action="store_true", default=False,
                      dest="export")

other_group.add_option("--export-format",
                      help="Format for exports ('task' or 'batch')",
                      action="store", default="task",
                      dest="export_format")

other_group.add_option("--format",
                      help="Output format: plain (default), html, json or none",
                      action="store", type="string",
                      default=None,
                      dest="format")

other_group.add_option("--import",
                      help="Read JSON task template from a file, - for stdin",
                      action="store", type="string",
                      dest="importfile")

other_group.add_option("--output",
                      help="Write result(s) to a file, substituting a run number for %n",
                      action="store", type="string",
                      default=None,
                      dest="output")

other_group.add_option("--full",
                       help="Produce full diagnostics and archiving information with result",
                       action="store_true", default=False,
                       dest="full")

other_group.add_option("--quiet",
                      help="Operate quietly",
                      action="store_true", default=False,
                      dest="quiet")

other_group.add_option("--lead-bind",
                       help="Have the lead participant Bind to the supplied address for administrative communications",
                       action="store", type="string",
                       default=None,
                       dest="lead_bind")

other_group.add_option("--url",
                      help="Dump a URL that points to the task after posting and exit",
                      action="store_true", default=False,
                      dest="url")

other_group.add_option("--debug",
                       help="Debug this (task) command",
                       action="store_true",
                       dest="debug")

other_group.add_option("--debug-pscheduler",
                       help="Enable pScheduler debug for this task",
                       action="store_true",
                       dest="debug_pscheduler")






(options, remaining_args) = opt_parser.parse_args()

if len(remaining_args) < 1 and options.importfile is None:
    opt_parser.print_usage()
    pscheduler.fail()


task_schema = pscheduler.HighInteger(1)
schedule_schema = pscheduler.HighInteger(1)

#
# Validate the command line
#

if options.max_runs < 1:
    pscheduler.fail("Invalid --max-runs; must be 1 or more")


formats = {
    'html': 'text/html',
    'json': 'application/json',
    'none': None,
    'text': 'text/plain',
    # Not "officially" supported, but here for completeness
    'text/html': 'text/html',
    'application/json': 'application/json',
    'text/plain': 'text/plain',
    }

try:
    opt_format = "text/plain" if options.format is None else options.format
    out_format = formats[opt_format]
except KeyError:
    pscheduler.fail("Invalid --format; must be text, html, json or none")


if options.repeat is not None:
    try:
        repeat = pscheduler.iso8601_as_timedelta(options.repeat)
    except ValueError as ex:
        pscheduler.fail("Invalid --repeat: %s" % str(ex))


if options.repeat_cron is not None:
    schedule_schema.set(2)
    repeat_cron = options.repeat_cron

if options.key is not None:
    try:
        key = pscheduler.string_from_file(options.key)
    except IOError as ex:
        pscheduler.fail("Unable to read key file: " + str(ex))


if options.reference is not None:
    try:
        reference = pscheduler.string_from_file(options.reference)
        reference = pscheduler.json_load(reference)
    except IOError as ex:
        pscheduler.fail("Unable to read reference file: " + str(ex))
    except ValueError as ex:
        pscheduler.fail("Invalid --reference '%s': %s"
                        % (options.reference, str(ex)))


    

if options.start is not None:
    # TODO: Support "Pxx" and "@Pxx" formats like the database does
    # Should have a module function that does this.
    try:
        start = pscheduler.iso8601_as_datetime(options.start, localize=True)
    except ValueError as ex:
        pscheduler.fail("Invalid --start: %s" % str(ex))
    if start <= pscheduler.time_now():
        pscheduler.fail("Invalid --start; must be in the future.")
    pass

if options.until is not None:
    # TODO: Support "Pxx" and "@Pxx" formats like the database does
    # TODO: Error handling in iso8601_as_datetime() needs improvement.
    try:
        until = pscheduler.iso8601_as_datetime(options.until, localize=True)
    except ValueError as ex:
        pscheduler.fail("Invalid --until: %s" % str(ex))
    if until <= pscheduler.time_now():
        pscheduler.fail("Invalid --until; must be in the future.")


verbose = (not options.quiet) \
    and not (options.url) \
    and options.format is None

# TODO: Tie this to the options.
log = pscheduler.Log(verbose=verbose, debug=options.debug, quiet=True, propagate=True)


# Decide who assists and where to bind for it.

last_ditch_assist = pscheduler.api_local_host()
assist = options.assist or os.getenv('PSCHEDULER_ASSIST') or last_ditch_assist
assist_bind = options.assist_bind or os.getenv('PSCHEDULER_ASSIST_BIND')
log.debug("Assistance is from %s bound from %s", assist, assist_bind)


#
# If we were asked to read in some JSON, do that.  Anything the
# options add will override it.
#

if options.importfile is None:
    task = {
        'schedule': {},
        'test': {
            'spec': {}
            }
        }
else:
    if options.importfile == '-':
        file = sys.stdin
    else:
        try:
            file = open(options.importfile)
        except IOError as ex:
            pscheduler.fail("Unable to open task %s" % (str(ex)))

    try:
        task = pscheduler.json_load(file, exit_on_error=True)
    except ValueError as ex:
        pscheduler.fail("Invalid task specification: %s" % str(ex))

    # Validate the JSON against a TaskSpecification
    # TODO: Figure out how to do this without the intermediate object
    valid, message = pscheduler.json_validate({"": task}, {
        "type": "object",
        "properties": {
            "": {"$ref": "#/pScheduler/TaskSpecification"}
        },
        "required": [""]
    })

    if not valid:
        pscheduler.fail("Invalid imported JSON: %s" % (message))

    # An empty schedule is needed later for handling the slip time.
    if 'schedule' not in task:
        task['schedule'] = {}



if options.slip is not None:

    # The command-line argument overrides everything.
    log.debug("Using slip from command line")
    task['schedule']['slip'] = options.slip

else:

    # If there's no slip already in the schedule and we're not
    # importing or exporting, force a default.

    if ('slip' not in task['schedule']) \
       and (not options.export) \
       and (options.importfile is None):
        forced_slip = os.environ.get('PSCHEDULER_SLIP', 'PT5M')
        task['schedule']['slip'] = forced_slip
        log.debug("Forcing default slip of %s" % (forced_slip))



# Overlay the lead bind if specified
if options.lead_bind is not None:
    task["lead-bind"] = options.lead_bind

# Overlay the key if specified
if options.key is not None:
    task["_key"] = key

# Overlay the reference if specified
if options.reference is not None:
    task["reference"] = reference

#
# Overlay schedule options
#

# Put a default empty schedule in the task which will be removed on
# export if empty.
if 'schedule' not in task:
    task['schedule'] = {}

if options.max_runs > 1:
    task['schedule']['max-runs'] = options.max_runs

if options.priority is not None:
    task['priority'] = options.priority
else:
    try:
        task['priority'] = int(os.environ['PSCHEDULER_PRIORITY'])
    except ValueError:
        pscheduler.fail("PSCHEDULER_PRIORITY does not contain an integer.")
    except KeyError:
        pass


if 'priority' in task:
    task_schema.set(3)


if options.repeat is not None:
    task['schedule']['repeat'] = options.repeat

if options.repeat_cron is not None:
    task['schedule']['repeat-cron'] = options.repeat_cron

if options.sliprand is not None:
    task['schedule']['sliprand'] = options.sliprand

if options.start is not None:
    task['schedule']['start'] = options.start

if options.until is not None:
    task['schedule']['until'] = options.until


if options.debug_pscheduler:
    task_schema.set(4)
    task['debug'] = True




bind = options.bind
lead_bind = options.lead_bind


#
# Figure out what kind of test this is.  Don't worry about it being
# valid, that will be checked later.
#

if len(remaining_args) > 0:
    test_arg = remaining_args.pop(0)
    if test_arg != '-':
        task['test']['type'] = test_arg

test_type = task['test'].get('type', '-')

if test_type == '-':
    pscheduler.fail("No test type specified.")


# Add desired tools, if any.

if options.tool:
    task['tools'] = options.tool


# Add archivers, if any.
if options.archive:

    archives = []

    for archive in options.archive:
        try:
            archive_text = pscheduler.string_from_file(archive)
            archive_json = pscheduler.json_load(archive_text)
            archives.append(archive_json)
        except IOError as ex:
            pscheduler.fail("Unable to read archive file: %s" % (str(ex)))
        except ValueError as ex:
            pscheduler.fail("Archiver '%s': %s" % (archive, str(ex)))

    task['archives'] = archives


if options.keep_after_archive:
    task['keep-after-archive'] = options.keep_after_archive
    task_schema.set(5)

# Add contexts, if any.
if options.context:

    try:
        context_text = pscheduler.string_from_file(options.context)
        context_json = pscheduler.json_load(context_text, max_schema=1)
    except IOError as ex:
        pscheduler.fail("Unable to read context file: %s" % (str(ex)))
    except ValueError as ex:
        pscheduler.fail("Context '%s': %s" % (context_text, str(ex)))

    task['contexts'] = context_json

# If there's anything context-related (explicit or imported), bump the
# schema to a version that supports it.
if 'contexts' in task:
    task_schema.set(2)



# Make sure whoever we're using for assistance is running pScheduler
(has, error) = pscheduler.api_has_pscheduler(assist, bind=assist_bind)
if not has:
    pscheduler.fail("Assist server %s is unavailable: %s" % (assist, error))


#
# Convert the remaining arguments to a test spec.
#

# Resolve any task arguments beginning with @.

try:
    task_args = [ pscheduler.string_from_file(arg) for arg in remaining_args ]
except Exception as ex:
    pscheduler.fail(str(ex))


spec_url = pscheduler.api_url_hostport(assist,
                              path='/tests/' + test_type + '/spec')
log.debug("Converting to spec via %s", spec_url)
status, raw_spec = pscheduler.url_get(
    spec_url,
    params={ 'args': pscheduler.json_dump(task_args) },
    bind=assist_bind,
    throw=False,
    timeout=10
    )

if status == 400:

    # Anything with --help or -h in it is a plea for help.
    if "--help" in remaining_args or "-h" in remaining_args:
        pscheduler.succeed(
            "Usage: task [task-options] %s [test-options]\n\n%s"
            % (test_type, raw_spec))

    # Anything else is a bona-fide bad request.

    pscheduler.fail("%s: %s" % (assist, raw_spec))

if status == 404:
    pscheduler.fail("Could not find test " + test_type + " on server")

if status == 500:
    pscheduler.fail("Internal error on on %s.  Consult system logs for details." \
                        % ( "local pScheduler server" if assist == last_ditch_assist
                            else assist ))
elif status != 200:
    pscheduler.fail("Unknown error %d: %s" % (status, raw_spec))

json_to_merge = raw_spec
assert 'spec' in task['test']
final_schema = max(task['test']['spec'].get('schema', 1),
                   json_to_merge.get("schema", 1))
task['test']['spec'].update(json_to_merge)
if final_schema > 1:
    task['test']['spec']['schema'] = final_schema


if task_schema.value() > 1:
    task['schema'] = task_schema.value()

if schedule_schema.value() > 1:
    task['schedule']['schema'] = schedule_schema.value()


# Export

if options.export:

    # If the schedule item ended up empty, get rid of it.
    try:
        if len(task['schedule']) == 0:
            del task['schedule']
    except KeyError:
        pass

    if options.export_format == 'task':
        to_export = task
    elif options.export_format == 'batch':
        try:
            del task['schedule']
        except KeyError:
            pass  # Not there?  Don't care.

        to_export = {
            'schema': 1,
            'jobs': [
                {
                    'label': 'Exported %s' % (test_arg),
                    'task': task
                }
            ]
        }
    else:
        pscheduler.fail("Invalid export format '%s'" % (options.export_format))

    pscheduler.json_dump(obj=to_export, dest=sys.stdout, pretty=True)
    print()
    pscheduler.succeed()


#
# Contact the assist server
#

if verbose:
    print(("Submitting with assistance from %s..." % assist) if assist != last_ditch_assist \
        else "Submitting task...")



# TODO: Validate the test before figuring out who's involved.

#
# Determine the lead participant
#

lead = None
null_reason = None

# For the time being, try the later version of the API, falling back
# on the older one.  This should only be a problem for a few days.
# TODO: Remove fallback prior to 4.0; see #53.

# Get participants

url = pscheduler.api_url_hostport(
    assist, '/tests/%s/participants' % task['test']['type'])
log.debug("Fetching participant list")

spec_text = pscheduler.json_dump(task['test']['spec'])
participants_params = {'spec': spec_text}
log.debug("Spec is: %s", spec_text)
log.debug("Params are: %s", participants_params)
status, participants = pscheduler.url_get(
    url,
    params=participants_params,
    bind=assist_bind,
    throw=False )

if status not in [ 200, 404 ]:
    pscheduler.fail("Unable to determine the lead participant: %s" % participants)
if status == 200:
    log.debug("Got participants: %s", participants)
    lead = participants["participants"][0]
    null_reason = participants.get("null-reason")

if null_reason is None:
    null_reason = "No reason provided."

log.debug("Lead is %s", lead)
if lead is None:
    log.debug("Null reason is %s", null_reason)

# If the lead is None, the usual behavior would be to task the server
# on the local host.  If an assist server is being used, that's
# probably a good indication that the local system has no server.
# Barf mightily.

if lead is None and assist != last_ditch_assist:
    pscheduler.fail("Cannot use an assist server when the test parameters are ambiguous.\n" \
                    "A parameter like --host, --host-node, --source or --source-node that\n" \
                    "provides a hint about where the test is to run will prevent this error.")



#
# Check that the lead is running pScheduler and is reasonably responsive.
#


if lead is None:
    lead = pscheduler.api_local_host()

log.debug("Pinging lead %s", lead)
(up, reason) = pscheduler.api_ping(lead)
if not up:
    pscheduler.fail("Unable to find pScheduler on %s: %s." % (lead, reason))


#
# Give the task to the lead for scheduling.
#

tasks_url = pscheduler.api_url(lead, '/tasks')
log.debug("Posting task to %s", tasks_url)
log.debug("Data is %s", task)
status, task_url = pscheduler.url_post(
    tasks_url,
    data=task,
    bind=bind,
    throw=False)

if status != 200:
    pscheduler.fail("Failed to post task: " + task_url +
        "\nThe 'pscheduler troubleshoot' command may be of use in problem" + 
        "\ndiagnosis. 'pscheduler troubleshoot --help' for more information.")

# If asked to just dump the URL, do that and exit.
if options.url:
    pscheduler.succeed(task_url)

if verbose:
    print("Task URL:")
    print(task_url)
log.debug("Posted %s", task_url)

#
# Spit out the tool being used
#

status, result = pscheduler.url_get(
    task_url,
    bind=bind,
    params={ "detail": 1 },
    json=True,
    throw=False)

if status != 200:
    pscheduler.fail("Failed to fetch the task: %s" % (result))

try:
    log.debug("Submission diagnostics:\n%s" \
              % (pscheduler.indent(result['detail']['diags'])))
except KeyError:
    log.debug("No submission diagnostics available.")

if verbose:
    print("Running with tool '%s'" % (result['tool']))


#
# Get the first future run.
#

if verbose:
    print("Fetching first run...")

# TODO: It would be more RESTful to have this URL available as part of
# the task instead of building it here.

runs_url = task_url + '/runs/first'
log.debug("Fetching %s", runs_url)
status, run_json = pscheduler.url_get(runs_url, bind=bind, throw=False)

if status == 404:
    pscheduler.fail("%s never scheduled a run for the task." % (
        "The local host" if lead is None else lead))

# Watch the task run.


watch_args = [ "pscheduler", "watch",
               "--first"
               ]

# Force a format only if explicitly told to do so.
if options.format is not None:
    watch_args.extend(["--format", out_format])

if options.debug:
    watch_args.append("--debug")
if options.full:
    watch_args.append("--full")
if options.quiet:
    watch_args.append("--quiet")
if options.bind is not None:
    watch_args.append("--bind")
    watch_args.append(options.bind)
if options.output is not None:
    watch_args.append("--output")
    watch_args.append(options.output)
# TODO: Pass --verbose once watch supports that.
watch_args.append(task_url)

watch_args = " ".join([ shlex.quote(arg) for arg in watch_args ])
log.debug("Handing off: %s", watch_args)
os.execl("/bin/sh", "/bin/sh", "-c", watch_args)
