#!/usr/bin/python
#
# pScheduler Archiver Daemon
#


import collections
import datetime
import optparse
import os
import signal
import threading
import time

import daemon
import pscheduler
import psycopg2
import psycopg2.extensions

from dateutil.tz import tzlocal


pscheduler.set_graceful_exit()



# 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

opt_parser.add_option("-a", "--archive-defaults",
                      help="Directory containing default archivers",
                      action="store", type="string", dest="archive_defaults",
                      default="/etc/pscheduler/default-archives")
opt_parser.add_option("-c", "--channel",
                      help="Schedule notification channel",
                      action="store", type="string", dest="channel",
                      default="archiving_change")
opt_parser.add_option("-d", "--dsn",
                      help="Database connection string",
                      action="store", type="string", dest="dsn",
                      default="dbname=pscheduler")
opt_parser.add_option("-m", "--max-parallel",
                      help="Maximum concurrent archivings",
                      action="store", type="int", dest="max_parallel",
                      default=50)
opt_parser.add_option("-p", "--pool-size",
                      help="Size of pool per archive type",
                      action="store", type="int", dest="pool_size",
                      default=15)
opt_parser.add_option("-r", "--refresh",
                      help="Forced refresh interval (ISO8601)",
                      action="store", type="string", dest="refresh",
                      default="PT15S")
opt_parser.add_option("-s", "--skim-interval",
                      help="How often to skim archiver pools (ISO8601)",
                      action="store", type="string", dest="skim_interval",
                      default="PT30S")
opt_parser.add_option("--timeout",
                      help="Timeout for archiver I/O (ISO8601)",
                      action="store", type="string", dest="timeout",
                      default="PT2M")
opt_parser.add_option("--retry",
                      help="Retry interval after I/O timeout (ISO8601)",
                      action="store", type="string", dest="retry",
                      default="PT30S")
opt_parser.add_option("--verbose", action="store_true", dest="verbose")
opt_parser.add_option("--debug", action="store_true", dest="debug")


(options, args) = opt_parser.parse_args()

if options.max_parallel < 1:
    opt_parser.error("Number of concurrent archivings must be positive.")

if options.pool_size < 1:
    opt_parser.error("Pool size must b be positive.")

refresh = pscheduler.iso8601_as_timedelta(options.refresh)
if refresh is None:
    opt_parser.error('Invalid refresh interval "' + options.refresh + '"')
refresh = pscheduler.timedelta_as_seconds(refresh)



skim_interval = pscheduler.iso8601_as_timedelta(options.skim_interval)
if skim_interval is None:
    opt_parser.error('Invalid skim interval "' + options.skim_interval + '"')
skim_interval = pscheduler.timedelta_as_seconds(skim_interval)


timeout = pscheduler.iso8601_as_timedelta(options.timeout)
if timeout is None:
    opt_parser.error('Invalid timeout "' + options.timeout + '"')
timeout = pscheduler.timedelta_as_seconds(timeout)

retry = pscheduler.iso8601_as_timedelta(options.retry)
if retry is None:
    opt_parser.error('Invalid retry "' + options.retry + '"')
retry = pscheduler.timedelta_as_seconds(retry)


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

dsn = options.dsn


#
# Maintainer for default archives
#

class DefaultArchiveMaintainer:

    def __init__(self, path, dsn, log):

        self.path = path
        self.dsn = dsn
        self.log = log

        # This needs its own database connection without autocommit.
        self.db = pscheduler.PgConnection(dsn, name="archiver-default",
                                          autocommit=False)

        # Most recent file or directory change we saw
        self.most_recent = 0


    def refresh(self):

        self.log.debug("Refreshing default archivers from %s", self.path)

        if not os.path.isdir(self.path):
            self.log.debug("%s is not a directory", self.path)
            return

        timestamps = [ os.path.getmtime(self.path) ]

        # Examine everything before tinkering with the database

        archives = []

        # Hold warnings until after we see something's changed so we
        # don't spew them repeatedly.
        warnings = []

        try:
            paths = [ os.path.join(self.path, f)
                      for f in os.listdir(self.path) ]
        except OSError as ex:
            self.log.debug("Unable to read %s: %s", self.path, ex)
            return

        # Gather the timestamps

        for path in paths:

            self.log.debug("Examining file %s", path)

            try:

                if os.path.isfile(path):

                    # Append the timestamp whether the file is valid or
                    # not, because we're looking for any kind of change.
                    # Note that we check mtime *and* ctime to catch
                    # permission changes and the like.
                    timestamps.append( max(os.path.getctime(path),
                                           os.path.getmtime(path)) )
                else:
                    self.log.debug("Not a file")
                    continue

            except Exception as ex:
                warnings.append("Ignoring %s: %s" % (path, str(ex)))
                continue

        # If nothing's changed, we're done.

        timestamps = sorted(timestamps, reverse=True)
        newest = timestamps[0]
        if newest <= self.most_recent:
            self.log.debug("Nothing has changed; not updating.")
            return

        self.most_recent = newest

        # Parse the files and make a list to be inserted into the database

        for path in paths:

            self.log.debug("Reading file %s", path)

            try:
                with open(path, 'r') as content:
                    spec = pscheduler.json_load(content, max_schema=1)
            except Exception as ex:
                warnings.append("Ignoring %s: %s" % (path, str(ex)))
                continue

            archives.append({
                "path": path,
                "spec": spec
            })

        # Spit out any warnings that were accumulated

        for warning in warnings:
            self.log.warning(warning)

        # Write to the database.  Note that we don't take any steps to
        # suppress exceptions; if something goes south we want the
        # entire program to restart.

        self.log.debug("Writing the database")

        try:
            self.db.query("DELETE FROM archive_default")
        except Exception as ex:
            self.log.error("Failed to delete defaults: %s", str(ex))

        for archive in archives:

            self.db.query("SAVEPOINT archiver_insert")
            failed = True
            try:
                self.db.query("INSERT INTO archive_default (archive) VALUES (%s)",
                              [ pscheduler.json_dump(archive["spec"]) ])
                self.log.debug("Inserted data from %s", archive["path"])
                failed = False
            except psycopg2.Error as ex:
                self.log.warning("Ignoring %s: %s", archive["path"], ex.diag.message_primary)
            except Exception as ex:
                self.log.error("%s: %s", archive["path"], str(ex))

            if failed:
                self.db.query("ROLLBACK TO SAVEPOINT archiver_insert")
            else:
                self.db.query("RELEASE SAVEPOINT archiver_insert")

        # Finish up and get out

        self.db.commit()




#
# Archive Processs, ProcessPools and Pool Collections
#


class ArchiveProcess(object):
    """
    Maintain an archive process that streams JSON in and out.
    """

    def __init__(self, archiver):
        self.archiver = archiver
        self.program = None
        self.__establish()

    def __establish(self):
        """
        Start the archiver process
        """
        if self.program is None:
            self.program = pscheduler.StreamingJSONProgram([
                "pscheduler", "internal", "invoke",
                "archiver", self.archiver, "archive"
            ], timeout=timeout)

    def __call__(self, json):
        """
        Do a round trip with the program.
        """
        self.__establish()
        try:
            return self.program(json)
        except (IOError, pscheduler.StreamingJSONProgramFailure) as ex:
            self.done()
            return {
                "succeeded": False,
                "error": "Possibly-permanent archiver error: %s" % (str(ex))
                # Don't retry.  The StreamingJSONProgram will have retried.
                }


    def done(self):
        """Stop the archive process"""
        if self.program is not None:
            self.program.done()
            self.program = None


    def __del__(self):
        self.done()




class ArchiverProcessPool(object):
    """
    A pool of archivers of a single type
    """

    def __init__(self, archiver, log, max_size, skim_age=60):
        self.archiver = archiver
        self.log = log
        self.max_size = max_size
        self.skim_age = skim_age

        self.pool = collections.deque()
        self.lock = threading.Lock()
        self.size = 0
        self.reset_cycle = 0
        self.skim_count = 0
        self.next_skim = time.time() + skim_age
        self.max_utilization = 0


    def __len__(self):
        return self.size


    def __call__(self, json):
        """
        Pull a process from the pool (or create one) and run a result
        through it.
        """

        process = None

        while process is None:
            with self.lock:
                if len(self.pool) > 0:
                    # Take a process from the pool.
                    process = self.pool.pop()
                else:
                    # Pool is empty
                    if self.size < self.max_size:
                        # Room for new processes; create one.
                        process = ArchiveProcess(self.archiver)
                        self.size += 1
                        self.log.debug("Pool %s: New process %d",
                                       self.archiver, self.size)
                    else:
                        # Pool is full and all processes are busy.
                        pass

                if process is not None:
                    self.log.debug("Pool %s: Got a process", self.archiver)
                    cycle_at_start = self.reset_cycle
                    self.max_utilization = max(self.max_utilization,
                                               self.size - len(self.pool))

            # TODO: Need to do something better than this, but it has
            # to be deadlock-proof.
            if process is None:
                self.log.debug("Pool %s: Waiting for a process", self.archiver)
                time.sleep(0.25)


        try:
            result = process(json)
        finally:
            # Dispose of the process no matter what happened.
            with self.lock:
                same_cycle = self.reset_cycle == cycle_at_start
                if same_cycle \
                   and self.skim_count == 0 \
                   and self.size <= self.max_size:
                    self.pool.append(process)
                else:
                    process.done()
                    self.size -= 1
                    if same_cycle and self.skim_count > 0:
                        self.skim_count -= 1
                    self.log.debug("Pool %s: Dropped a process", self.archiver)

        return result


    def fill(self):
        """
        Artificially add processes up to the maximum.  (Intended for test
        and debug only.)
        """
        with self.lock:
            while self.size < self.max_size:
                self.pool.append(ArchiveProcess(self.archiver))
                self.size += 1
            self.log.debug("Pool %s: Filled to %d", self.archiver, self.size)


    def drain(self):
        """
        Force all running processes to stop.
        """
        with self.lock:
            # Anything in the pool is idle
            self.size -= len(self.pool)
            for proc in self.pool:
                proc.done()
            self.pool.clear()

            # Changing reset cycles makes runners not re-pool themselves
            self.reset_cycle += 1

        self.log.debug("Pool %s: Drained", self.archiver)



    def skim(self):
        """
        Clean out anything beyond peak utilization since the last skim.
        """
        now = time.time()
        if now < self.next_skim:
            return

        self.log.debug("Pool %s: Skimming", self.archiver)

        with self.lock:

            if self.max_utilization < self.size:
                to_remove = self.size - self.max_utilization 
                self.log.debug("Pool %s: Removing %d/%d processes for %d",
                               self.archiver, to_remove, self.size,
                               self.max_utilization)

                # Idle processes get removed first
                while to_remove > 0 and len(self.pool) >= to_remove:
                    self.pool.pop().done()
                    self.size -= 1
                    to_remove -= 1

                    # Anything left comes out of what's in circulation
                    self.skim_count = to_remove

            # Reset for next time.
            self.max_utilization = 0
            self.next_skim = now + self.skim_age


    def done(self):
        self.drain()


    def __del__(self):
        self.done()




class ArchiverPoolCollection(object):
    """
    A collection of ArchiverProcessPools for each archiver type.
    """

    def __init__(self, log, max_size=10, skim_age=60):
        self.log = log
        self.max_size = max_size
        self.skim_age = skim_age

        self.pool = {}
        self.lock = threading.Lock()


    def __call__(self, archiver, json):
        """
        Archive a result to the correct pool.
        """

        with self.lock:
            if archiver not in self.pool:
                self.pool[archiver] = ArchiverProcessPool(
                    archiver, self.log,
                    max_size=self.max_size, skim_age=self.skim_age)

        assert archiver in self.pool
        return self.pool[archiver](json)


    def skim(self):
        """
        Skim all pools
        """
        with self.lock:
            del_list = []
            for archiver_name in self.pool:
                archiver = self.pool[archiver_name]
                archiver.skim()


    def drain(self):
        """
        Drain all pools
        """
        with self.lock:
            for archiver in self.pool:
                self.pool[archiver].drain()


    def done(self):
        self.drain()





# Dictionary of archivings in progress

# TODO: Use the ThreadSafeSet and ThreadSafeSetHold instead
workers = pscheduler.ThreadSafeDictionary()


#
# Archive Worker
#

class ArchiveWorker():

    def __init__(self, db, log, row, collection):
        self.db = db
        self.log = log
        self.row = row
        self.collection = collection

        self.id = row[0]

        self.worker = threading.Thread(target=lambda: self.run())
        self.worker.setDaemon(True)
        self.worker.start()


    def run(self):
        """
        Archive the result in a thread-safe way
        """
        self.log.debug("%d: Thread running", self.id)
        try:
            self.__run()
        except Exception as ex:
            # Don't worry about the result here.  If __run() failed to
            # post anything, that will be the end of it.  If it did,
            # it might be salvageable.
            self.log.exception()
        self.log.debug("%d: Thread finished", self.id)
        del workers[self.id]


    def __run(self):
        """
        Do the deed
        """

        failed = False

        archiving_id, task_uuid, run_uuid, archiver, archiver_data, start, \
            duration, test, tool, participants, result_merged, attempts, \
            last_attempt, transform, task_detail, run_detail, spec = self.row

        participants_merged = []
        for participant in participants:
            participants_merged.append(pscheduler.api_local_host() \
                                       if participant is None \
                                       else participant)

        uri_host = spec.get("uri-host", pscheduler.api_local_host_fqdn())
        task_href = pscheduler.api_url(host=uri_host, path="tasks/%s" % (task_uuid))
        run_href = pscheduler.api_url(host=uri_host, path="tasks/%s/runs/%s" \
                                      % (task_uuid, run_uuid))

        # Strip the internal-use-only parts from the task details.
        for remove in [
                'archives'
        ]:
            try:
                del task_detail[remove]
            except KeyError:
                pass


        # Strip the internal-use-only parts from the run details.
        for remove in [
                'archivings',
                'participant-data',
                'participant-data-full',
                'result'
        ]:
            try:
                del run_detail[remove]
            except KeyError:
                pass

        self.log.debug("%d: Task is %s", self.id, task_href)
        self.log.debug("%d: Run is %s", self.id, run_href)

        json = {
            # This may contain private data that the archiver needs to see.
            'data': archiver_data,
            'task-href': task_href,
            'run-href': run_href,
            'result': {
                'id': run_uuid,
                'schedule': {
                    'start': pscheduler.datetime_as_iso8601(start),
                    'duration': pscheduler.timedelta_as_iso8601(duration)
                },
                # This may contain data that isn't to be leaked.
                'test': pscheduler.json_decomment(test, prefix="_", null=True),
                'tool': {
                    'name': tool['name'],
                    'version': tool['version'],
                },
                'run': run_detail,
                'task': task_detail,
                'participants': participants_merged,
                'result': result_merged,
                'reference': task_detail.get('reference', None)
            },
            'attempts': attempts,
            'last-attempt': None if last_attempt is None \
            else pscheduler.datetime_as_iso8601(last_attempt)
        }

        json_result = json["result"]

        json_result["task"]["href"] = task_href
        json_result["run"]["href"] = run_href
        json_result["run"]["task-href"] = task_href

        # If there's a transform (already validated), do it.

        if transform is not None:
            self.log.debug("%d Transforming input %s", self.id, json)
            self.log.debug("%d Script is %s", self.id, transform["script"])
            try:
                raw = transform.get("output-raw", False)
                transformer = pscheduler.JQFilter(
                    filter_spec=transform.get("script", "."),
                    args=transform.get("args", {}),
                    output_raw=raw
                )
                if raw:
                    json["result"] = transformer(json["result"])
                else:
                    # Result is always a single object when doing JSON.
                    json["result"] = transformer(json["result"])[0]
                self.log.debug("%d: Transformed to %s", self.id, json)
            except Exception as ex:
                self.log.error("%d: Error during transformation: %s", str(ex))
                failed = True
                returncode = 1
                stdout = ""
                stdout_attempt = ""
                stderr = "Error during transformation: %s" % (str(ex))
                result = {
                    "succeeded": False,
                    "error": stderr
                }


        if json["result"] is None:
            self.log.debug("%d: Null transform result; not archiving", self.id)
            returncode = 0
            stdout = """{ "succeeded": true }"""
            stdout_attempt = "Transform result forced not archiving."
            stderr = ""
            skip = True
            result = {
                "succeeded": True,
                "error": ""
            }
        else:
            skip = False

        if (not failed) and (not skip):

            self.log.debug("%d: Archiving to %s: %s",
                           self.id, archiver, json)

            try:
                result = self.collection(archiver, json)
                self.log.debug("%d: Returned JSON from archiver: %s", self.id, result)
                returncode = 0
            except Exception as ex:
                result = {
                    "succeeded": False,
                    "error": "Exception during archiving: %s" % (str(ex))
                }
                self.log.exeption("%d: Exception" % (self.id))
                returncode = 1


        attempt = {
            # TODO: Figure out to format this with -xx:xx for the timezone offset.
            "time": pscheduler.datetime_as_iso8601(datetime.datetime.now(tzlocal())),
            "return-code": returncode,
            "stdout": result,
            "stderr": result.get("error", "")
        }

        if returncode != 0:

            self.log.debug("%d: Permanent Failure: %s", self.id,
                           result.get("error", "Unspecified error"))
            self.db.query("""UPDATE archiving
                             SET
                                 completed = TRUE,
                                 attempts = attempts + 1,
                                 last_attempt = now(),
                                 next_attempt = NULL,
                                 diags = diags || (%s::JSONB)
                             WHERE id = %s""",
                          [ pscheduler.json_dump(attempt), archiving_id ])

        else:

            if result['succeeded']:
                self.log.debug("%d: Succeeded: %s to %s",
                               self.id, run_uuid, archiver)
                self.db.query("""UPDATE archiving
                                 SET
                                     completed = TRUE,
                                     archived = TRUE,
                                     attempts = attempts + 1,
                                     last_attempt = now(),
                                     next_attempt = NULL,
                                     diags = diags || (%s::JSONB)
                                 WHERE id = %s""",
                              [ pscheduler.json_dump(attempt), archiving_id ])

            else:

                self.log.warning("%d: Failed to archive %s to %s: %s",
                                 self.id, run_href, archiver,
                                 result.get("error", "Unspecified problem"))

                # If there's a retry, schedule the next one.

                if "retry" in result:

                    next_delta = pscheduler.iso8601_as_timedelta(
                        result['retry'])

                    next = datetime.datetime.now(tzlocal()) \
                           + next_delta

                    self.log.debug("%d: Rescheduling for %s", self.id, next)

                    self.db.query("""UPDATE archiving
                                     SET
                                         attempts = attempts + 1,
                                         last_attempt = now(),
                                         next_attempt = %s,
                                         diags = diags || (%s::JSONB)
                                     WHERE id = %s""",
                                  [next,
                                   pscheduler.json_dump(attempt),
                                   archiving_id])

                else:

                    self.log.debug("%d: No retry requested.  Giving up.",
                                   self.id)
                    self.log.warning("%d: Gave up archiving %s to %s",
                                     self.id, run_href, archiver)


                    self.db.query("""UPDATE archiving
                                     SET
                                         completed = TRUE,
                                         attempts = attempts + 1,
                                         last_attempt = now(),
                                         next_attempt = NULL,
                                         diags = diags || (%s::JSONB)
                                     WHERE id = %s""",
                                  [pscheduler.json_dump(attempt), archiving_id])








#
# Main Program
#

def main_program():

    db = pscheduler.PgConnection(dsn)
    db.listen(options.channel)
    db.listen("warmboot")

    archiver_collection = ArchiverPoolCollection(
        max_size=options.pool_size,
        skim_age=skim_interval,
        log=log)


    # Something to maintain the default archiver list
    default_maintainer = DefaultArchiveMaintainer(options.archive_defaults, dsn, log)
    default_maintainer.refresh()

    next_refresh = None
    last_debug = log.is_forced_debugging()

    db.query("SELECT heartbeat_boot('archiver')")


    while True:

        db.query("SELECT heartbeat('archiver')")

        # Drain the pools if the debug level changes.  This is done
        # since persisting archivers won't register the change in state.

        current_debug = log.is_forced_debugging()
        if current_debug != last_debug:
            archiver_collection.drain()
            last_debug = current_debug
        else:
            archiver_collection.skim()


        warmboot = False

        # Wait for something to happen.

        if next_refresh is None:

            log.debug("Retrieving immediately.")

        else:

            next_refresh = min(next_refresh, refresh)

            log.debug("Waiting %s for change or notification", next_refresh)

            warmboot = False

            db.query("SELECT heartbeat('archiver', %s)", [datetime.timedelta(seconds=next_refresh)])

            if db.wait(next_refresh):
                notifications = [ tup[0] for tup in db.notifications() ]
                log.debug("Notifications: %s" % (" ".join(notifications)))
                warmboot = "warmboot" in notifications
            else:
                log.debug("Nothing happened.")

        # Until we hear otherwise...
        next_refresh = refresh

        if warmboot:
            log.debug("Warm boot; draining archiver pools")
            archiver_collection.drain()

        # If we're already full up on workers, don't bother hitting the database
        if len(workers) == options.max_parallel:
            log.debug("Already have a full slate of workers.")
            continue

        # TODO: A lot of what's being pulled here can be pulled from
        # the spec.  Remove anything unnecessary.
        result = db.query("""SELECT id, task_uuid, run_uuid, archiver,
                             archiver_data, start,
                             duration, test, tool, participants, result,
                             attempts, last_attempt, transform, task_detail, run_detail,
                             spec
                             FROM archiving_next(%s)""",
                          [options.max_parallel])

        if len(result) == 0:
            log.debug("Nothing to archive; finding time until next archiving.")
            result = db.query("""SELECT next_attempt - now() FROM archiving
                                 WHERE NOT completed AND next_attempt > now()
                                 ORDER BY next_attempt LIMIT 1""")

            assert len(result) < 2
            if len(result) == 0:
                log.debug("Nothing in the future.")
            else:
                next_refresh = pscheduler.timedelta_as_seconds(result.next()[0])
                log.debug("Next archiving in %s", next_refresh)

            continue

        # Only refresh the default archivers after a wait that got rows.
        if next_refresh is not None:
            default_maintainer.refresh()

        log.debug("Got %d rows", len(result))

        for row in result:

            # Don't bother if there are already too many archivers running.
            if len(workers) >= options.max_parallel:
                log.debug("Already running %d archivers.", len(workers))
                break

            id = row[0]

            if id in workers:
                log.debug("%d: Already running a worker", id)
                continue

            log.debug("%d: Starting worker", id)
            workers[id] = ArchiveWorker(db, log, row, archiver_collection)


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())
