"""Zookeeper Partitioner Implementation

:class:`SetPartitioner` implements a partitioning scheme using
Zookeeper for dividing up resources amongst members of a party.

This is useful when there is a set of resources that should only be
accessed by a single process at a time that multiple processes
across a cluster might want to divide up.

Example Use-Case
----------------

- Multiple workers across a cluster need to divide up a list of queues
  so that no two workers own the same queue.

"""
import logging
import os
import socket
import time
import threading
from functools import partial

from kazoo.client import KazooState


log = logging.getLogger(__name__)


class PartitionState(object):
    """High level partition state values

    .. attribute:: ALLOCATING

        The set needs to be partitioned, and may require an existing
        partition set to be released before acquiring a new partition
        of the set.

    .. attribute:: ACQUIRED

        The set has been partitioned and acquired.

    .. attribute:: RELEASE

        The set needs to be repartitioned, and the current partitions
        must be released before a new allocation can be made.

    .. attribute:: FAILURE

        The set partition has failed. This occurs when the maximum
        time to partition the set is exceeded or the Zookeeper session is
        lost. The partitioner is unusable after this state and must
        be recreated.

    """
    ALLOCATING = "ALLOCATING"
    ACQUIRED = "ACQUIRED"
    RELEASE = "RELEASE"
    FAILURE = "FAILURE"


class SetPartitioner(object):
    """Partitions a set amongst members of a party

    This class will partition a set amongst members of a party such
    that each member will be given zero or more items of the set and
    each set item will be given to a single member. When new members
    enter or leave the party, the set will be re-partitioned amongst the
    members.

    When the :class:`SetPartitioner` enters the
    :attr:`~PartitionState.FAILURE` state, it is unrecoverable
    and a new :class:`SetPartitioner` should be created.


    Simple Example:

    .. code-block:: python

        from kazoo.client import KazooClient
        client = KazooClient()

        qp = client.SetPartitioner(
            path='/work_queues', set=('queue-1', 'queue-2', 'queue-3'))

        def use_setlist(partitions):
            # do something with a partition passed in

        # Run the use_setlist function when partitions are acquired
        # repeatedly
        qp.run(func=use_setlist)

    Sometimes, more control is needed over handling the state
    transitions. Or the program may need to do other things during
    specific states, in which case a more verbose example that allows
    for finer grained control is needed:

    .. code-block:: python

        from kazoo.client import KazooClient
        client = KazooClient()

        qp = client.SetPartitioner(
            path='/work_queues', set=('queue-1', 'queue-2', 'queue-3'))

        while 1:
            if qp.failed:
                raise Exception("Lost or unable to acquire partition")
            elif qp.release:
                qp.release_set()
            elif qp.acquired:
                for partition in qp:
                    # Do something with each partition
            elif qp.allocating:
                qp.wait_for_acquire()

    **State Transitions**

    When created, the :class:`SetPartitioner` enters the
    :attr:`PartitionState.ALLOCATING` state.

    :attr:`~PartitionState.ALLOCATING` ->
    :attr:`~PartitionState.ACQUIRED`

        Set was partitioned successfully, the partition list assigned
        is accessible via list/iter methods or calling list() on the
        :class:`SetPartitioner` instance.

    :attr:`~PartitionState.ALLOCATING` ->
    :attr:`~PartitionState.FAILURE`

        Allocating the set failed either due to a Zookeeper session
        expiration, or failure to acquire the items of the set within
        the timeout period.

    :attr:`~PartitionState.ACQUIRED` ->
    :attr:`~PartitionState.RELEASE`

        The members of the party have changed, and the set needs to be
        repartitioned. :meth:`SetPartitioner.release` should be called
        as soon as possible.

    :attr:`~PartitionState.ACQUIRED` ->
    :attr:`~PartitionState.FAILURE`

        The current partition was lost due to a Zookeeper session
        expiration.

    :attr:`~PartitionState.RELEASE` ->
    :attr:`~PartitionState.ALLOCATING`

        The current partition was released and is being re-allocated.

    """
    def __init__(self, client, path, set, partition_func=None,
                 identifier=None, time_boundary=30):
        """Create a :class:`~SetPartitioner` instance

        :param client: A :class:`~kazoo.client.KazooClient` instance
        :param path: The partition path to use
        :param set: The set of items to partition
        :param partition_func: A function to use to decide how to
                               partition the set
        :param identifier: An identifier to use for this member of the
                           party when participating. Defaults to the
                           hostname + process id.
        :param time_boundary: How long the party members must be stable
                              before allocation can complete

        """
        self.state = PartitionState.ALLOCATING

        self._client = client
        self._path = path
        self._set = set
        self._partition_set = []
        self._partition_func = partition_func or self._partitioner
        self._identifier = identifier or '%s-%s' % (
            socket.getfqdn(), os.getpid())
        self._locks = []
        self._lock_path = '/'.join([path, 'locks'])
        self._party_path = '/'.join([path, 'party'])
        self._time_boundary = time_boundary

        self._acquire_event = client._handler.event_object()

        # Create basic path nodes
        client.ensure_path(path)
        client.ensure_path(self._lock_path)
        client.ensure_path(self._party_path)

        # Join the party
        self._party = client.ShallowParty(self._party_path,
                                          identifier=self._identifier)
        self._party.join()

        self._state_change = threading.Condition()
        client.add_listener(self._establish_sessionwatch)

        # Now watch the party and set the callback on the async result
        # so we know when we're ready
        self._children_updated = False
        self._child_watching(self._allocate_transition, async=True)

    def __iter__(self):
        """Return the partitions in this partition set"""
        for partition in self._partition_set:
            yield partition

    @property
    def failed(self):
        """Corresponds to the :attr:`PartitionState.FAILURE` state"""
        return self.state == PartitionState.FAILURE

    @property
    def release(self):
        """Corresponds to the :attr:`PartitionState.RELEASE` state"""
        return self.state == PartitionState.RELEASE

    @property
    def allocating(self):
        """Corresponds to the :attr:`PartitionState.ALLOCATING`
        state"""
        return self.state == PartitionState.ALLOCATING

    @property
    def acquired(self):
        """Corresponds to the :attr:`PartitionState.ACQUIRED` state"""
        return self.state == PartitionState.ACQUIRED

    def wait_for_acquire(self, timeout=30):
        """Wait for the set to be partitioned and acquired

        :param timeout: How long to wait before returning

        """
        self._acquire_event.wait(timeout)

    def release_set(self):
        """Call to release the set

        This method begins the step of allocating once the set has
        been released.

        """
        self._release_locks()
        if self._locks:
            # This shouldn't happen, it means we couldn't release our
            # locks, abort
            self._fail_out()
            return
        else:
            with self._state_change:
                if self.failed:
                    return
                self.state = PartitionState.ALLOCATING
        self._child_watching(self._allocate_transition, async=True)

    def finish(self):
        """Call to release the set and leave the party"""
        self._release_locks()
        self._fail_out()

    def _fail_out(self):
        with self._state_change:
            self.state = PartitionState.FAILURE
        if self._party.participating:
            try:
                self._party.leave()
            except Exception:
                pass

    def _allocate_transition(self, result):
        """Called when in allocating mode, and the children settled"""
        # Did we get an exception waiting for children to settle?
        if result.exception:
            self._fail_out()
            return

        children, async_result = result.get()
        self._children_updated = False

        # Add a callback when children change on the async_result
        def updated(result):
            with self._state_change:
                if self.acquired:
                    self.state = PartitionState.RELEASE
            self._children_updated = True

        async_result.rawlink(updated)

        # Split up the set
        self._partition_set = self._partition_func(
            self._identifier, list(self._party), self._set)

        # Proceed to acquire locks for the working set as needed
        for member in self._partition_set:
            if self._children_updated or self.failed:
                # Still haven't settled down, release locks acquired
                # so far and go back
                return self._abort_lock_acquisition()

            lock = self._client.Lock(self._lock_path + '/' +
                                     str(member))
            try:
                lock.acquire()
            except Exception:
                return self.finish()
            self._locks.append(lock)

        # All locks acquired! Time for state transition, make sure
        # we didn't inadvertently get lost thus far
        with self._state_change:
            if self.failed:
                return self.finish()
            self.state = PartitionState.ACQUIRED
            self._acquire_event.set()

    def _release_locks(self):
        """Attempt to completely remove all the locks"""
        self._acquire_event.clear()
        for lock in self._locks[:]:
            try:
                lock.release()
            except Exception:
                # We proceed to remove as many as possible, and leave
                # the ones we couldn't remove
                pass
            else:
                self._locks.remove(lock)

    def _abort_lock_acquisition(self):
        """Called during lock acquisition if a party change occurs"""
        self._partition_set = None
        self._release_locks()
        if self._locks:
            # This shouldn't happen, it means we couldn't release our
            # locks, abort
            self._fail_out()
            return
        return self._child_watching(self._allocate_transition)

    def _child_watching(self, func=None, async=False):
        """Called when children are being watched to stabilize

        This actually returns immediately, child watcher spins up a
        new thread/greenlet and waits for it to stabilize before
        any callbacks might run.

        """
        watcher = ChildrenWatcher(self._client, self._party_path,
                                  self._time_boundary)
        asy = watcher.start()
        if func:
            # We spin up the function in a separate thread/greenlet
            # to ensure that the rawlink's it might use won't be
            # blocked
            if async:
                func = partial(self._client._handler.spawn, func)
            asy.rawlink(func)
        return asy

    def _establish_sessionwatch(self, state):
        """Register ourself to listen for session events, we shut down
        if we become lost"""
        if state == KazooState.LOST:
            self._client.remove_listener(self._establish_sessionwatch)
            self._fail_out()

    def _partitioner(self, identifier, members, partitions):
        # Ensure consistent order of partitions/members
        all_partitions = sorted(partitions)
        workers = sorted(members)

        i = workers.index(identifier)
        # Now return the partition list starting at our location and
        # skipping the other workers
        return all_partitions[i::len(workers)]


class ChildrenWatcher(object):
    """Node Children Watcher that alerts when the children stabilize

    A separate watcher for the children of a node, that ignores
    changes within a boundary time and sets the result only when the
    boundary time has elapsed and there is no change of children for
    the boundary time's duration.

    Example::

        watcher = ChildrenWatcher(client, '/some/path')
        result = watcher.start()

        # Blocks until the children have not changed for time boundary
        # seconds, returns list of children and an async_result that
        # will be set once if the children change in the future
        children, async_result = result.get()

    """
    def __init__(self, client, path, time_boundary=30):
        self.client = client
        self.path = path
        self.time_boundary = time_boundary
        self.children = []
        self.last_change = None

    def start(self):
        """Begin the watching process asynchronously

        :returns: An :class:`~kazoo.interfaces.IAsyncResult` instance
                  that will be set when no change has occurred to the
                  children for time boundary seconds.

        """
        self.asy = asy = self.client._handler.async_result()
        self.client._handler.spawn(self._inner_start)
        return asy

    def _inner_start(self):
        if not self.last_change:
            self.last_change = time.time()

        try:
            async_result = None
            while time.time() - self.last_change < self.time_boundary:
                async_result = self.client._handler.async_result()
                self.children = self.client.retry(
                    self.client.get_children, self.path,
                    partial(self.children_watcher, async_result))

                # Now sleep for a time boundary's worth from the last change
                sleep_till = self.time_boundary - (time.time() -
                                                   self.last_change)
                time.sleep(sleep_till)

            self.asy.set((self.children, async_result))
        except Exception as exc:
            self.asy.set_exception(exc)

    def children_watcher(self, async, event):
        self.last_change = time.time()
        async.set(time.time())
