#!/usr/bin/env python3
#
# Dump the contents of the schedule in a plottable form
#

# General syntax:
#     pscheduler plot-schedule [options]

import datetime
import optparse
import os
import pscheduler
import pytz
import shutil
import sys
import tempfile
import urllib.parse


pscheduler.set_graceful_exit()


#
# Utilities
#

def get_time_with_delta(string):
    """
    Get an absolute time or delta and return an ISO8601 string with
    the absolute time.
    """

    # If it looks like an ISO time, return that.
    try:
        absolute = pscheduler.iso8601_as_datetime(string)
        # Default behavior is to localize naive times.
        if absolute.tzinfo is None:
            absolute = pytz.utc.localize(absolute)
            return pscheduler.datetime_as_iso8601(absolute)
    except ValueError:
        pass

    try:
        if string[0:1] == "+P":
            delta = pscheduler.iso8601_as_timedelta(string[1:])
        elif string[0:1] == "-P":
            delta = -1 * pscheduler.iso8601_as_timedelta(string[1:])
        else:
            pass
    except ValueError:
        pscheduler.fail("Invalid time delta '%s'" % (string))

    # Let this throw what it's going to throw.
    delta = pscheduler.iso8601_as_timedelta(string)

    return pscheduler.datetime_as_iso8601(
        pscheduler.time_now() + delta)


#
# Gargle the arguments
#

whoami = os.path.basename(sys.argv[0])
args = sys.argv[1:]


# Pre-convert any trailing arguments that look like times or deltas
# into raw times so the option parser doesn't choke on negative
# deltas.

if len(args) > 0:
    arg = -1
    while True and abs(arg) <= len(args):
        try:
            args[arg] = (get_time_with_delta(str(args[arg])))
        except ValueError:
            break
        arg -= 1



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

opt_parser = VerbatimParser(
    usage="Usage: %prog [ OPTIONS ] [ delta | start end ] ",
    epilog=
"""
Example:

  plot-schedule
      Plot the schedule on the local host for the next hour

  plot-schedule --host ps3.example.net
      Plot the schedule on ps3.example.net for the next hour

  plot-schedule -PT1H
      Plot the schedule on the local host for an hour in the
      the past

  plot-schedule +PT25M
      Plot the schedule on the local host for 25 minutes in
      the future

  schedule -PT1H +PT30M
      Plot the schedule on the local host for an hour in the
      the past and 30 minutes into the future

  schedule 2016-05-01T12:40:00 2016-05-01T12:55:00
      Plot the schedule on the local host between the times
      specified.
"""
    )
opt_parser.disable_interspersed_args()

opt_parser.add_option("--bind",
                      help="Make the request from the provided address",
                      default=None,
                      action="store", type="string",
                      dest="bind")

opt_parser.add_option("--host",
                      help="Request schedule from named host",
                      default=pscheduler.api_local_host(),
                      action="store", type="string",
                      dest="host")

opt_parser.add_option("--plot",
                      help="Write GNUPlot commands to standard output and exit",
                      action="store_true",
                      dest="plot")



(options, remaining_args) = opt_parser.parse_args(args)



full_host = options.host

# This outputs PNG, which isn't suitable for TTYs.

if sys.stdout.isatty() and not options.plot:
    pscheduler.fail("Not sending PNG output to a tty.  Consider redirecting.")



now = pscheduler.time_now()

if len(remaining_args) == 0:

    # Default; show an hour's worth.
    start = now
    end = start + datetime.timedelta(hours=1)

elif len(remaining_args) == 1:

    # One argument is an absolute time or a timedelta.

    try:
        arg = pscheduler.iso8601_as_datetime(remaining_args[0])
    except ValueError:
        pscheduler.fail("Invalid time specification")

    if arg < now:
        start = arg
        end = now
    else:
        start = now
        end = arg

elif len(remaining_args) == 2:

    try:
        start = pscheduler.iso8601_as_datetime(remaining_args[0])
        end = pscheduler.iso8601_as_datetime(remaining_args[1])
    except ValueError:
        pscheduler.fail("Invalid time specification")


    if end < start:
        start, end = end, start

else:
    pscheduler.fail(usage)


# Find the server's time zone and convert all times to that.

status, clock = pscheduler.url_get(
    pscheduler.api_url(host=options.host, path="clock"),
    bind=options.bind,
    throw = False
    )


if status != 200:
    pscheduler.fail("Failed to get server's clock:  %d: %s" % (status, clock))

try:
    server_clock = clock["time"]
except KeyError:
    pscheduler.fail("Server did not return its clock.")

server_time = pscheduler.iso8601_as_datetime(server_clock)
server_timezone = server_time.tzinfo

# Convert the start and end times to the server's time zone

# Pylint thinks these are ints.
# pylint: disable=maybe-no-member
start = start.astimezone(server_timezone)
end = end.astimezone(server_timezone)

#
# If the server is API 5 or higher, let it do the work.
#
status, api = pscheduler.url_get(
    pscheduler.api_url(host=options.host, path="api"),
    bind=options.bind,
    throw = False
    )

if status != 200:
    pscheduler.fail("Failed to get server's API level:  %d: %s" % (status, api))

try:
    api_level = int(api)
except ValueError:
    pscheduler.fail("Server returned a malformed API level.")

if api_level >= 5:

    schedule_url = pscheduler.api_url(host=options.host or None, path="/schedule")
    query = urllib.parse.urlencode({
        "start": start.isoformat(),
        "end": end.isoformat(),
        "format": "png"
    })
    curl_args = [ "curl", "-s", "-k", "-A", "pScheduler plot-schedule" ]
    if options.bind is not None:
        curl_args.extend(["--interface", options.bind])
    curl_args.append("{}?{}".format(schedule_url, query))

    #print(curl_args)
    os.execvp(curl_args[0], curl_args)
    #pscheduler.fail("Not implemented yet.")
    #exit(99)


#
# Locally-generated graph for older servers
#

# TODO: This can be removed once 4.3.x is no longer supported.


#
# Fetch the schedule
#

status, schedule = pscheduler.url_get(
    pscheduler.api_url(host=options.host, path="schedule"),
    params={
        "start": pscheduler.datetime_as_iso8601(start),
        "end": pscheduler.datetime_as_iso8601(end)
        },
    bind=options.bind,
    throw = False
    )

if status != 200:
    pscheduler.fail("Server returned status %d: %s" % (status, schedule))


# Where we place each class and The list of columns depends on whether
# or not the background classes are combined.

class_levels = {
    "exclusive": 1,
    "normal": 2,
    "background": 3,
    "background-multi": 4,
    "nonstart": 5,
    "preempted": 6
}

# TODO: These could be pulled from the database, but there would need
# to be an API endpoint for it.
columns = [
    "\\n",
    "\\nExclusive",
    "\\nNormal",
    "Background\\nSingle-Result",
    "Background\\nMulti-Result",
    "Non\\nStarting",
    "\\nPreempted"
    ]


script_terminal = [
"""
reset
set terminal png notransparent truecolor size 800,1200 background rgb "#ffffff"
"""
]

script_lines = [
"""
set timefmt "%Y-%m-%dT%H:%M:%S"

unset xtics
""",
"set x2range [0.5:{0}.5]".format(len(columns) - 1),
"""
set x2tics out scale 0 ( \\\\
""",
", ".join([ '"{0}" {1}'.format(pair[1], pair[0])
        for pair in enumerate(columns, start=0) ]) + "\\\\",
"""
    )


set ydata time
set ytics out nomirror
set format y "%Y-%m-%d\\n%H:%M:%S"
""",
"""set yrange ["%s":"%s"] reverse""" \
    % (pscheduler.datetime_as_iso8601(start),
       pscheduler.datetime_as_iso8601(end)),
"""set title "pScheduler Schedule for %s\\n%s""" % (
    full_host, server_time.strftime("%F %T%z")),
"""
set key off
set grid front noxtics ytics mytics linetype 0
set boxwidth 0.9

set style fill solid border lc rgb "#00e000"

plot "-" using 1:2:3:2:3 \\
  with candlesticks \\
  linewidth 1 \\
  linecolor rgb "#00e000" \\
  axes x2y1"""
]

#
# Dump it out
#

if len(schedule) == 0:
    pscheduler.fail("Nothing scheduled %s to %s" % (
        pscheduler.datetime_as_iso8601(start),
        pscheduler.datetime_as_iso8601(end)
        ))


for run in schedule:


    # Force times to server's time zone
    for field in [ "start-time", "end-time" ]:
        in_time = pscheduler.iso8601_as_datetime(run[field])
        converted_time = in_time.astimezone(server_timezone)
        run[field] = pscheduler.datetime_as_iso8601(converted_time)

    try:
        run_state = run["state"]
        sched_class = run_state if run_state in [ "nonstart" ] \
                      else run["test"]["scheduling-class"]
    except KeyError:
        pscheduler.fail("%s: Software too old to support this feature" \
                        % (host))

    script_lines.append(
        "%s %s %s" % (
            class_levels[sched_class],
            run["start-time"],
            run["end-time"]
            )
    )

script_lines.append("eof\n")


#
# Write the plot file if that's what we were asked to do.
#

if options.plot:
    print("\n".join(script_terminal),
          "\n".join(script_lines)
    )
    pscheduler.succeed()


#
# Push it through GNUPlot to make a PNG
#

# run_program() can no longer gargle non-text output from programs.
# Make a temp file, send the plot there and do a faux cat to get it to
# stdout.

png = None
try:
    (png_file, png_path) = tempfile.mkstemp()
    os.close(png_file)

    script = "\n".join(script_terminal) \
             + "\nset output '%s'\n" % (png_path) \
             + "\n".join(script_lines)

    status, out, err = pscheduler.run_program("gnuplot", stdin=script)

    if status != 0:
        pscheduler.fail(err)

    with open(png_path, "rb") as image:
        shutil.copyfileobj(image, sys.stdout.buffer)

#except Exception as ex:
#    pscheduler.fail("Failed to generate plot: %s" % (ex))

finally:
    if png_path is not None:
        os.unlink(png_path)

if sys.stderr.isatty():
    sys.stderr.write("Plotted %d runs\n" % (len(schedule)))

pscheduler.succeed()
