#
# __init__.py
#
# This source file is part of the FoundationDB open source project
#
# Copyright 2013-2018 Apple Inc. and the FoundationDB project authors
# 
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# 
#     http://www.apache.org/licenses/LICENSE-2.0
# 
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

# FoundationDB TaskBucket layer

import random, uuid, time, struct
import fdb, fdb.tuple
fdb.api_version(200)

# TODO: Make this fdb.tuple.subspace() or similar?
class Subspace (object):
    def __init__(self, prefixTuple, rawPrefix=""):
        self.rawPrefix = rawPrefix + fdb.tuple.pack(prefixTuple)
    def __getitem__(self, name):
        return Subspace( (name,), self.rawPrefix )
    def key(self):
        return self.rawPrefix
    def pack(self, tuple):
        return self.rawPrefix + fdb.tuple.pack( tuple )
    def unpack(self, key):
        assert key.startswith(self.rawPrefix)
        return fdb.tuple.unpack(key[len(self.rawPrefix):])
    def range(self, tuple=()):
        p = fdb.tuple.range( tuple )
        return slice(self.rawPrefix + p.start, self.rawPrefix + p.stop)

def random_key():
    return uuid.uuid4().bytes

def _pack_value(v):
    if hasattr(v, 'pack'):
        return v.pack()
    else:
        return v

class TaskTimedOutException(Exception):
    pass

class TaskBucket (object):
    """A TaskBucket represents an unordered collection of tasks, stored
    in a database and available to any number of clients.  A program may
    use one or many TaskBuckets.  TaskBucket represents a task as a
    key/value dictionary.  See TaskDispatcher for an easy way to define
    tasks as functions."""

    def __init__(self, subspace, system_access = False):
        self.prefix = subspace
        self.active = self.prefix["ac"]
        self.available = self.prefix["av"]
        self.timeouts = self.prefix["to"]
        self.timeout = 30000
        self.system_access = system_access

    @fdb.transactional
    def clear(self, tr):
        """Removes all tasks, whether or not locked, from the bucket."""
        if self.system_access:
            tr.options.set_access_system_keys()
        del tr[ self.prefix.range(()) ]

    @fdb.transactional
    def add(self, tr, taskDict):
        """Adds the task specified by taskDict to the bucket for any participant to do."""
        if self.system_access:
            tr.options.set_access_system_keys()
        assert taskDict
        key = random_key()
        for k,v in taskDict.items():
            tr[ self.available.pack( (key, k) ) ] = _pack_value( v )
        taskDict["__task_key"] = key
        return key

    @fdb.transactional
    def addIdle(self, tr):
        """Adds an idle task for any participant to remove."""
        if self.system_access:
            tr.options.set_access_system_keys()
        key = random_key()
        tr[ self.available.pack( (key, "type") ) ] = ""
        return key

    @fdb.transactional
    def get_one(self, tr):
        """Gets a single task from the bucket, locks it so that only the
        caller will work on it for a while, and returns its taskDict.
        If there are no tasks in the bucket, returns None."""
        if self.system_access:
            tr.options.set_access_system_keys()
        k = tr.snapshot.get_key( fdb.KeySelector.last_less_or_equal( self.available.pack( (random_key(),) ) ) )
        if not k or k < self.available.pack( ("",) ):
            k = tr.snapshot.get_key( fdb.KeySelector.last_less_or_equal( self.available.pack( (chr(255)*16,) ) ) )
            if not k or k < self.available.pack( ("",) ):
                if self.check_timeouts(tr):
                    return self.get_one(tr)
                return None
        key = self.available.unpack(k)[0]
        avail = self.available[key]
        timeout = tr.get_read_version().wait() + long( self.timeout * (0.9 + 0.2*random.random()) )

        taskDict = {}
        for k,v in tr[ avail.range(()) ]:
            tk, = avail.unpack(k)
            taskDict[tk]=v

            if tk != "type" or v != "":
                tr[ self.timeouts.pack( (timeout, key, tk) ) ] = v
        del tr[ avail.range(()) ]
        tr[self.active.key()] = random_key()

        taskDict["__task_key"] = key
        taskDict["__task_timeout"] = timeout
        return taskDict

    @fdb.transactional
    def is_empty(self, tr):
        if self.system_access:
            tr.options.set_read_system_keys()
        k = tr.get_key( fdb.KeySelector.last_less_or_equal( self.available.pack( (chr(255)*16,) ) ) )
        if k and k >= self.available.pack( ("",) ):
            return False
        return not bool(next(iter(tr[self.timeouts.range()]),False))

    @fdb.transactional
    def is_busy(self, tr):
        if self.system_access:
            tr.options.set_read_system_keys()
        k = tr.get_key( fdb.KeySelector.last_less_or_equal( self.available.pack( (chr(255)*16,) ) ) )
        return k and k >= self.available.pack( ("",) )

    @fdb.transactional
    def finish(self, tr, taskDict):
        """taskDict is a task previously locked by get_one.  Removes it permanently
        from the bucket.  If the task has already timed out, raises TaskTimedOutException."""
        if self.system_access:
            tr.options.set_access_system_keys()
        rng = self.timeouts.range( (taskDict["__task_timeout"], taskDict["__task_key"]) )
        if next(iter(tr[rng]),False):
            del tr[ rng ]
        else:
            raise TaskTimedOutException()

    @fdb.transactional
    def is_finished(self, tr, taskDict):
        #print "checking if the task was finished at version: {0}".format(tr.get_read_version().wait())
        if self.system_access:
            tr.options.set_read_system_keys()
        rng = self.timeouts.range( (taskDict["__task_timeout"], taskDict["__task_key"]) )
        return not bool(next(iter(tr[rng]),False))

    def check_active(self, db):
        @fdb.transactional
        def get_starting_value(tr):
            if self.system_access:
                tr.options.set_read_system_keys()
            if not self.is_busy(tr):
                self.addIdle(tr)
            return tr[self.active.key()]

        starting_value = get_starting_value(db)

        @fdb.transactional
        def get_active_key(tr):
            if self.system_access:
                tr.options.set_read_system_keys()
            if tr[self.active.key()] != starting_value:
                return True
            return False

        for i in range(10):
            time.sleep(0.5)
            if get_active_key(db):
                return True
        return False

    @fdb.transactional
    def extend(self, tr, taskDict):
        """Extends the lock on a task previously locked by get_one for another self.timeout seconds.
        If the task has already timed out, raises TaskTimedOutException."""
        pass

    def check_timeouts(self, tr):
        """Looks for tasks that have timed out and returns them to be available tasks.  Returns True
        iff any tasks were affected."""
        end = tr.get_read_version().wait()
        rng = slice( self.timeouts.range((0,)).start , self.timeouts.range((end,)).stop )
        anyTimeouts = False
        for k,v in tr.get_range( rng.start, rng.stop, streaming_mode = fdb.StreamingMode.want_all):
            timeout, taskKey, param = self.timeouts.unpack(k)
            anyTimeouts = True
            tr.set( self.available.pack( (taskKey,param) ), v )
        del tr[rng]

        return anyTimeouts

class TaskDispatcher (object):
    def __init__(self):
        self.taskTypes = {}

    def taskType(self, f):
        """Given a function, defines a task type with the same name as the
        function that calls the function.  The function should accept the
        task dictionary as keyword parameters.

        May be used as a decorator."""
        self.taskTypes[f.__name__] = f
        return f

    def makeTask(self, f, **kw):
        """Return a taskDict suitable for TaskBucket.add() based on a function
        which has already been passed to self.taskType() and keyword parameters
        for it."""
        assert self.taskTypes.get(f.__name__) == f
        assert 'type' not in kw
        d = dict(kw)
        d["type"] = f.__name__
        return d

    def dispatch(self, taskDict):
        """Calls the function associated with the type of the taskDict passed
        in, passing the taskDict as keyword arguments.  Does not finish or
        extend the task or otherwise interact with a TaskBucket."""
        if taskDict["type"] != "":
            self.taskTypes[ taskDict["type"] ]( **taskDict )

    def do_one(self, db, taskBucket):
        """Gets one task (if any) from the task bucket, executes it, and finishes it.
        Returns True if a task was executed, False if none was available."""
        task = taskBucket.get_one(db)
        if not task: return False
        self.dispatch(task)
        return True

class FutureBucket (object):
    """A factory for Futures, and a location in the database to store them."""
    def __init__(self, subspace, system_access = False):
        self.prefix = subspace
        self.dispatcher = TaskDispatcher()
        self.dispatcher.taskType( self._add_task )
        self.dispatcher.taskType( self._unblock_future )
        self.system_access = system_access

    @fdb.transactional
    def future(self, tr):
        """Returns a new Future which is unset."""
        if self.system_access:
            tr.options.set_access_system_keys()
        f = Future(self)
        f._add_block(tr, "")
        return f

    def unpack(self, packed_future):
        """Returns a Future such that Future.pack()==packed_future."""
        return Future( self, packed_future )

    @fdb.transactional
    def join(self, tr, *futures):
        """Returns a Future which is set when all of the given futures are set."""
        if self.system_access:
            tr.options.set_access_system_keys()
        joined = Future(self)
        joined._join(tr, futures)
        return joined

    def _add_task(self, tr, taskType, taskBucket, **task):
        bucket = TaskBucket( Subspace((), rawPrefix=taskBucket), system_access=self.system_access )
        task["type"] = taskType
        bucket.add( tr, task )

    def _unblock_future(self, tr, future, blockid, **task):
        if self.system_access:
            tr.options.set_access_system_keys()
        future = self.unpack(future)
        del tr[ future.prefix["bl"][blockid].key() ]
        if future.is_set(tr):
            future.perform_all_actions(tr)

class Future (object):
    """Represents a state which will become true ("set") at some point, and a set
    of actions to perform as soon as the state is true.  The state can change
    only once."""

    def __init__(self, bucket, key=None):
        """Not for direct use.  Call FutureBucket.future() instead."""
        if key is None: key = random_key()
        self.key = key
        self.bucket = bucket
        self.prefix = bucket.prefix[ self.key ]
        self.dispatcher = bucket.dispatcher
        self.system_access = bucket.system_access

    def pack(self):
        return self.key

    @fdb.transactional
    def is_set(self, tr):
        if self.system_access:
            tr.options.set_read_system_keys()
        return not any( tr[self.prefix["bl"].range(())] )

    @fdb.transactional
    def on_set_add_task(self, tr, taskBucket, taskDict):
        """When the Future is set, the given task will be added to the given taskBucket."""
        td = dict(taskDict)
        td["taskBucket"] = taskBucket.prefix.key()
        td["taskType"] = td["type"]
        td["type"] = "_add_task"
        self.on_set(tr, td)

    @fdb.transactional
    def on_set(self, tr, taskDict):
        """When the Future is set, the given task will be executed immediately
        as part of the same transaction.  Consider using onSetAddTask to defer
        the execution of the task instead."""
        if self.system_access:
            tr.options.set_access_system_keys()
        if self.is_set(tr):
            self.perform_action(tr, taskDict)
        else:
            cb_key = random_key()
            for k,v in taskDict.items():
                tr[ self.prefix["cb"][ cb_key ][k].key() ] = _pack_value( v )

    @fdb.transactional
    def set(self, tr):
        """Sets the Future immediately if it has not already been set.  Any
        actions that are to take place when it is set take place."""
        if self.system_access:
            tr.options.set_access_system_keys()
        #if self.is_set(tr): return
        del tr[ self.prefix.range( ("bl",) ) ]  # Remove all blocks
        self.perform_all_actions(tr)

    @fdb.transactional
    def join(self, tr, *futures):
        """This Future will be set when all of the given futures have been set.
        If this is called more than once on a Future, it is cumulative - the
        Future will be set when all of the Futures passed to any .join() are set."""
        if self.system_access:
            tr.options.set_access_system_keys()
        assert futures
        if self.is_set(tr): return
        del tr[ self.prefix["bl"][""].key() ]
        self._join(tr, futures)

    @fdb.transactional
    def joined_future(self, tr):
        """Returns a new unset future which this future is join()ed to."""
        if self.system_access:
            tr.options.set_access_system_keys()
        f = self.bucket.future(tr)
        self.join(tr, f)
        return f

    def _join(self, tr, futures):
        ids = [random_key() for f in futures]
        for blockid in ids:
            self._add_block(tr, blockid)
        for f,blockid in zip(futures,ids):
            f.on_set( tr, self.dispatcher.makeTask( self.bucket._unblock_future, future=self.pack(), blockid=blockid ) )

    def _add_block(self, tr, blockid):
        tr[ self.prefix["bl"][blockid].key() ] = ""

    def perform_all_actions(self, tr):
        cb = self.prefix["cb"]
        callbacks = list(tr[cb.range()])
        del tr[cb.range()] # Remove all callbacks
        # Now actually perform the callbacks
        taskDict = {}
        taskKey = None
        for k,v in callbacks:
            cb_key, k = cb.unpack(k)
            if cb_key != taskKey:
                self.perform_action( tr, taskDict )
                taskDict = {}
                taskKey = cb_key
            taskDict[k] = v
        self.perform_action(tr, taskDict)

    def perform_action(self, tr, taskDict):
        if taskDict:
            taskDict["tr"] = tr
            self.dispatcher.dispatch( taskDict )
