#!/usr/bin/python
#
# pScheduler Run Scheduler
#

import daemon
import datetime
import errno
import optparse
import pscheduler
import psycopg2
import psycopg2.extensions
import random
import re
import requests
import select
import socket
import sys
import time
import traceback
import urlparse


# Gargle the arguments

opt_parser = optparse.OptionParser()

# Daemon-related options

opt_parser.add_option("--daemon",
                      help="Daemonize",
                      action="store_true",
                      dest="daemon", default=False)
opt_parser.add_option("--pid-file",
                      help="Location of PID file",
                      action="store", type="string", dest="pidfile",
                      default=None)

# Program options

# TODO: Do we want pscheduler as the default here?
opt_parser.add_option("-d", "--dsn",
                      help="Database connection string",
                      action="store", type="string", dest="dsn",
                      default="dbname=pscheduler")
opt_parser.add_option("-r", "--refresh",
                      help="Forced refresh interval (ISO8601)",
                      action="store", type="string", dest="refresh",
                      default="PT10S")
opt_parser.add_option("-v", "--verbose", action="store_true", dest="verbose")
opt_parser.add_option("--debug", action="store_true", dest="debug")

(options, args) = opt_parser.parse_args()

refresh = pscheduler.iso8601_as_timedelta(options.refresh)
if refresh is None:
    opt_parser.error('Invalid refresh interval "' + options.refresh + '"')
if pscheduler.timedelta_as_seconds(refresh) == 0:
    opt_parser.error("Refresh interval must be calculable as seconds.")

log = pscheduler.Log(verbose=options.verbose, debug=options.debug)

dsn = options.dsn

# Minimum amount of time from now when the first run of a task can be
# scheduled.  This prevents "start now" tasks from being scheduled for
# a time before the participants can prepare for them.
# TODO: Potential race condition?  Yep.
first_run_offset = pscheduler.iso8601_as_timedelta('PT10S')



#
# Range Class and Functions
#

class Range():
    """
    Expresses ranges of values.  Values may be of any type with a
    less-than operator.
    """

    def __init__(self, lower, upper):
        self.lower = min(lower, upper)
        self.upper = max(lower, upper)

    def __repr__(self):
        return "R(%s..%s)" % (self.lower, self.upper)

    # Note that this is not __len__, which has to return an integer.
    def length(self):
        return self.upper - self.lower

    def __lt__(self, rhs):
        return self.lower < rhs.lower \
            or ( self.lower == rhs.lower and self.upper < rhs.upper)

    def __or__(self, rhs):
        """Find where two ranges overlap, return None if no overlap"""
        assert type(rhs) == type(self), "Wrong type"

        if self.lower > rhs.upper or self.upper < rhs.lower:
            return None
        return Range(max(self.lower, rhs.lower), min(self.upper, rhs.upper))

    def overlaps(self, candidates):
        """
        Find overlap in a list of candidate ranges, filtering out any that
        don't have any.
        """
        assert type(candidates) == list, "Wrong type"
        return [
            overlap for overlap in [
                self | x for x in candidates 
            ] if overlap is not None
        ]


def find_overlaps(range_lists):

    """
    Find a set of common ranges among several sets of them.

    (The use case for this in pScheduler is finding the common times
    that all participants in a task have available.)

    Applying this function to this input...

        [ [ Range(10, 50), Range(60, 77), Range(80, 90) ],
          [ Range(10, 12), Range(15,20), Range(20, 25), Range(51, 54),
            Range(65, 74), Range(81, 100) ],
          [ Range(22, 27), Range(65, 77), Range(82, 89) ],
          [ Range(1, 1000) ] ]

    ...will yield a result of [Range(22..25), Range(65..74), Range(82..89)].
    """

    sets = len(range_lists)

    # Common cases go first.

    if sets == 1:
        # One list merges with itself.
        return sorted(range_lists[0])

    if sets == 2:
        # Find overlaps in two lists

        # TODO: Figure out why we have to flatten this.
        def flatten(list_of_lists):
            return [ val for sublist in list_of_lists for val in sublist ]

        return sorted( flatten( [
            item for item in [
                item.overlaps(range_lists[1])
                for item in range_lists[0]
            ]
            if len(item) > 0
        ] ) )

    if sets == 0:
        # Nothing is... Nothing.
        return []

    # Anything more than two is recursive goodness to reduce the last
    # two sets into one and keep doing so until we get down to two.

    return sorted(find_overlaps(
        range_lists[0:-2] + [ find_overlaps(range_lists[-2:]) ]
    ))




#
# Run Poster
#

def run_post(
    url,           # URL for task
    start_time,    # Desired start time
    bind_addr,     # Bind address for HTTP or None
    log=None
    ):

    """
    Schedule a run of a task on all participating nodes.

    Returns a tuple containing:
        Run URI
        Scheduled start time
        Scheduled end time
        True if the run should be skipped and tried again later
        True if the above was because a scheduling conflict cropped up
        Error message
    """

    assert type(start_time) == datetime.datetime

    log and log.debug("Posting %s at %s", url, start_time)
    log and log.debug("Bind is %s", bind_addr)

    try:
        status, task = pscheduler.url_get(url, params={'detail': 1},
                                          bind=bind_addr, timeout=5)
    except pscheduler.URLException as ex:
        log and log.debug("Failed to fetch %s: %s", url, str(ex))
        return (None, None, None, False, False,
                "Error fetching task: %s" % (str(ex)))

    # Generate a list of the task URLs

    participants = task['detail']['participants']
    log and log.debug("Participant list is %s", participants)
    assert len(participants) >= 1

    # Make sure all of the participants' APIs answer before trying to
    # do anything with them.

    lead_bind = task.get("lead-bind", None)
    log and log.debug("Binding from %s", lead_bind)
    if not pscheduler.api_ping_all_up(participants, bind=lead_bind):

        # If the start time is after the next time we'd loop around,
        # we can try it again later.  Otherwise, tell the caller to
        # post a non-starter.

        skip = pscheduler.time_now() + refresh < start_time

        log and log.debug("Some participants down or slow. %s",
                          "Skipping." if skip else "Non-starter.") 
        return (None, None, None, skip, False,
                "One or more participants down or slow to respond.")
    else:
        log and log.debug("All participants are up.")

    task_urls = [ pscheduler.api_replace_host(url, participant)
                  for participant in participants ]
    log and log.debug("Task URLs are %s", task_urls)


    #
    # Figure out the range of times in which the task can be run.
    #

    task_duration = pscheduler.iso8601_as_timedelta(task['detail']['duration'])
    try:
        task_slip = pscheduler.iso8601_as_timedelta(task['schedule']['slip'])
    except KeyError:
        task_slip = datetime.timedelta()


    # If the task is a repeater, the run can't be slipped so far that
    # it would overlap with the next interval.  Adjust it accordingly.
    #
    # TODO: This should probably be enforced by the database when the
    # task is inserted.
    try:
        repeat_interval = pscheduler.iso8601_as_timedelta(task['schedule']['repeat'])
        if task_slip + task_duration >= repeat_interval:
            task_slip = repeat_interval - task_duration
            log and log.debug("Chopped slip to %s", task_slip)
    except KeyError:
        pass

    run_range_end = start_time + task_duration + task_slip

    range_params = {
        'start': pscheduler.datetime_as_iso8601(start_time),
        'end': pscheduler.datetime_as_iso8601(run_range_end)
        }

    log and log.debug("Possible run range %s", str(range_params))

    #
    # Get a list of the time ranges each participant has available to
    # run the task that overlap with the range we want.
    #

    log and log.debug("%s participants in this task", len(task_urls))

    # TODO: This would be good parallelized.
    # See http://stackoverflow.com/questions/5236364 for some ideas.

    range_set = []
    for task_url in task_urls:

        # TODO: It would be nice if the task had a list of the
        # runtimes URLs so we don't have to build it.
        runtime_url = task_url + '/runtimes'

        log and log.debug("Fetching proposals from %s", runtime_url)

        status, json_ranges = pscheduler.url_get(runtime_url,
                                                 params=range_params,
                                                 bind=bind_addr,
                                                 throw=False)

        if status in [ 404, 410 ]:
            log and log.debug("Task is no longer there.  Canceling %s.", url)
            del_status, del_result = pscheduler.url_delete(
                url, throw=False, bind=bind_addr)
            if del_status not in [ 200, 404, 410 ]:
                log and log.warning("Error while canceling %s: %s", url, del_result)
            return (None, None, None, False, False,
                    "Task is gone from %s" % (participant))

        if status != 200:
            log and log.debug("Got back %d: %s", status, json_ranges)
            return (None, None, None, False, False,
                    "Error trying to schedule with %s: %s %d"
                    % (participant, runtime_url, status))

        if len(json_ranges) == 0:
            log and log.debug("No time available.")
            return (None, None, None, False, False,
                    "%s has no time available for this run" % (participant))

        ranges = [ Range( pscheduler.iso8601_as_datetime(item['lower']),
                          pscheduler.iso8601_as_datetime(item['upper']) )
                   for item in json_ranges ]

        log and log.debug("Ranges: %s", ranges)
        
        range_set.append(ranges)

    log and log.debug("Done fetching time ranges")


    # Find the ranges all participants have in common

    common_ranges = [ trange for trange in find_overlaps(range_set)
                      if trange.length() >= task_duration ]
    log and log.debug("Ranges in common: %s", common_ranges)
    if len(common_ranges) == 0:
        return (None, None, None, False, False,
                "No common times for this run.")


    # If we're randomizing the start time, pick a range at random and
    # pick a random time within it.  Otherwise, take the earliest time
    # we can get.

    if ('sliprand' in task['schedule'] and task['schedule']['sliprand']):

        selected_range = random.choice(common_ranges)
        max_slip_offset = selected_range.length() - task_duration
        log and log.debug("Slipping randomly up to %s", max_slip_offset)

        if max_slip_offset:
            def us(td):
                return td.microseconds + 1000000 * (td.seconds + 86400 * td.days)
            increments = us(max_slip_offset) / 1000000
            picked_increment = random.randrange(0, increments)
            log and log.debug("%d increments, picked %d", increments, picked_increment)
        else:
            picked_increment = 0

        slip_offset = picked_increment * datetime.timedelta(seconds=1)

    else:

        log and log.debug("Taking earliest available time")
        selected_range = common_ranges[0]
        slip_offset = datetime.timedelta()

    if log:
        log.debug("Selected range %s", selected_range)
        log.debug("Using a slip offset of %s", slip_offset)
    schedule_lower = selected_range.lower + slip_offset
    schedule_upper = schedule_lower + task_duration

    if log:
        now = pscheduler.time_now()
        log.debug("Horizon range: %s - %s",
                  pscheduler.timedelta_as_iso8601(schedule_lower - now),
                  pscheduler.timedelta_as_iso8601(schedule_upper - now))



    #
    # Post the runs to the participants
    #

    run_params = { 'start-time': schedule_lower.isoformat() }

    runs_posted = []

    # First one is the lead.  Post it and get the UUID.

    if log:
        log.debug("Posting lead run to %s", task_urls[0])
        log.debug("Data %s", run_params)
    status, run_lead_url \
        = pscheduler.url_post(task_urls[0] + '/runs',
                              data=pscheduler.json_dump(run_params),
                              bind=bind_addr,
                              throw=False,
                              json=True)

    if status == 409:
        log and log.debug("Lead developed a schedule conflict.")
        return (None, None, None, True, True,
                "Lead developed a schedule conflict.")

    if status != 200:
        log and log.debug("Failed: %d %s", status, run_lead_url)
        return (None, None, None, False, False,
                "Failed to post lead run: %s" % run_lead_url)

    log and log.debug("Lead URL is %s", run_lead_url)
    assert type(run_lead_url) in [str, unicode]
    runs_posted.append(run_lead_url)

    # TODO: This should parse the URL and change the netloc instead of
    # assembling URLs.

    # What to add to a task URL to make the run URL
    run_suffix = run_lead_url[len(task_urls[0]):]

    # Cover the rest of the participants if there are any.

    errors = []
    conflict = False

    run_data = pscheduler.json_dump(run_params)

    for task_url in task_urls[1:]:

        put_url = task_url + run_suffix

        if log:
            log.debug("Putting run to participant %s", put_url)
            log.debug("Parameters: %s", run_params)

        status, output = pscheduler.url_put(put_url,
                                            data=run_data,
                                            bind=bind_addr,
                                            throw=False,
                                            json=False  # No output.
                                            )

        log and log.debug("PUT %d: %s", status, output)

        if status == 409:
            message = "%s developed a schedule conflict." % (
                urlparse.urlparse(request.url_root).netloc.split(':')[0])                
            log and log.debug(message)
            errors.append(message)
            conflict = True
            # No sense in continuing.
            break

        if status != 200:
            log and log.debug("Failed: %s", output)
            errors.append(output)
            continue

        runs_posted.append(put_url)
        log and log.debug("Succeeded.")

    if len(runs_posted) != len(task_urls):
        log and log.debug("Removing runs: %s", runs_posted)
        pscheduler.url_delete_list(runs_posted, bind=bind_addr)
        return (None, None, None, False, conflict,
                "Failed to post/put runs to all participants: %s"
                %  ("; ".join(errors))
        )

    #
    # Fetch all per-participant data, merge it and distribute the
    # result to all participants.
    #

    log and log.debug("Fetching per-participant data")

    part_data = []

    for run in runs_posted:

        # TODO: Should this be multiple attempts to avoid a race condition?
        log and log.debug("Getting part data from %s", run)
        status, result = pscheduler.url_get(run, bind=bind_addr, throw=False)
        if status != 200 or not 'participant-data' in result:
            log.debug("Deleting runs: %s", runs_posted)
            pscheduler.url_delete_list(runs_posted, bind=bind_addr)
            # TODO: Better error?
            return (None, None, None, False, False,
                    "Failed to get run data from all participants")
        part_data.append(result['participant-data'])
        log and log.debug("Got %s", result['participant-data'])

    full_data = pscheduler.json_dump ({
        'part-data-full': part_data
        })

    log and log.debug("Full part data: %s", full_data)

    for run in runs_posted:
        log and log.debug("Putting full part data to %s", run)
        status, result = pscheduler.url_put(run,
                                            data=full_data,
                                            bind=bind_addr,
                                            json=False,
                                            throw=False)
        if status != 200:
            pscheduler.url_delete_list(runs_posted, bind=bind_addr)
            # TODO: Better error?
            log and log.debug("Failed: %s", result)
            return (None, None, None, False, False,
                    "Failed to post run data to all participants")


    # TODO: Probably also want to return the start and end times?
    log and log.debug("Run posting finished")
    return (runs_posted[0], schedule_lower, schedule_upper, False, False, None)


#
# Main Program
#

def main_program():

    # TODO: All DB transactions need to be error checked

    pg = pscheduler.pg_connection(dsn)
    cursor = pg.cursor()
    cursor.execute("LISTEN task_change")

    # This is separate because inserts of non-starters happen while
    # fetching rows from the other cursor.
    nonstart_cursor = pg.cursor()


    while True:

        wait = True

        # TODO: The FALSE for background does nothing.  Get rid of it.
        cursor.execute("""
            SELECT uuid, runs, trynext, FALSE, json
            FROM schedule_runs_to_schedule LIMIT 1
            """)

        # Check if any notifications arrived while this query executed.
        if pg.notifies:
            wait = False
            del pg.notifies[:]

        # Any rows returned means we query again.
        if cursor.rowcount > 0:
            wait = False

        for row in cursor:

            uuid, runs, trynext, background, json = row

            log.debug("%sTASK %s, %d runs, try %s", 
                      "BACKGROUND " if background else "",
                      uuid, runs, trynext)

            # Punt the lead-bind value to "localhost" and make
            # everything would work nicely.  This is largely to avoid
            # the situation where the hostname points at an interface
            # that isn't up and is reasonably safe because we're only
            # talking to the lead, which is local.

            url = pscheduler.api_url(
                host=json.get("lead-bind", "localhost"),
                path="/tasks/%s" % (uuid))


            # For the first run only, push the start time out.
            # See comment above near the declaration of
            # first_run_offset.

            if runs == 0:
                later_start = pscheduler.time_now() + first_run_offset
                if trynext < later_start:
                    trynext = later_start                   

            log.debug("Trying to schedule %s for %s", uuid, trynext)
            log.debug("URL is %s", url)

            lead_bind = json.get("lead-bind", None)

            # Try a few times to schedule the run.  Sometimes,
            # something else may schedule a time we're working on, so
            # when that happens, try again.

            tries_left = 3
            while tries_left > 0:
                tries_left -= 1
                try:
                    run_uri, start_time, end_time, skip, conflict, error \
                        = run_post(url, trynext, lead_bind, log=log)
                except Exception as ex:
                    error = str(ex)
                    skip = False
                    break

                if not conflict:
                    break

                log.debug("Developed a scheduling conflict.  Trying again.")

            if tries_left == 0:
                skip = False
                error = "Gave up after too many scheduling conflicts."

            if skip:
                log.debug("Skipping: %s", error)
                continue

            if error is not None:
                log.debug("Unable: %s", error)
                # Post a non-starting run on the lead (local) system
                try:
                    nonstart_cursor.execute(
                        """SELECT api_run_post(%s, %s, NULL, %s)""",
                        [uuid, trynext, error])
                    log.info("Posted non-starting run at %s for %s: %s",
                             pscheduler.datetime_as_iso8601(trynext), url, error)
                except:
                    log.error("Insertion of non-starter failed: %s", error)

            else:
                log.debug("Scheduled for %s - %s at %s",
                          start_time, end_time, run_uri)


        # Wait for something to happen.
          
        if wait:

            try:
                if select.select([pg],[],[],
                                 pscheduler.timedelta_as_seconds(refresh)) \
                                 != ([],[],[]):
                    # Notified
                    pg.poll()
                    del pg.notifies[:]

            except select.error as ex:
                err_no, message = ex
                if err_no != errno.EINTR:
                    log.exception()
                    raise ex
                else:
                    log.debug("Interrupted during select")



if options.daemon:
    pidfile = pscheduler.PidFile(options.pidfile)
    with daemon.DaemonContext(pidfile=pidfile):
        pscheduler.safe_run(lambda: main_program())
else:
    pscheduler.safe_run(lambda: main_program())

