import logging
import threading

from kazoo.client import KazooClient
from kazoo.recipe.cache import TreeCache, TreeEvent


class TreeCacheListener(object):

    CONNECTIVE_EVENTS = {
        TreeEvent.CONNECTION_SUSPENDED: 'SUSPENDED',
        TreeEvent.CONNECTION_RECONNECTED: 'RECONNECTED',
        TreeEvent.CONNECTION_LOST: 'LOST',
    }

    def __init__(self, tree_cache):
        self._log = logging.getLogger('{0}.{1}'.format(
            __name__, self.__class__.__name__))
        self._cache = tree_cache
        self._initialized = threading.Event()

    def start(self, timeout=5):
        self._cache.listen(self._handle_event)
        self._cache.listen_fault(self._handle_fault)
        self._cache.start()
        return self._initialized.wait(timeout)

    def _handle_event(self, event):
        if not self._initialized.is_set():
            if event.event_type == TreeEvent.INITIALIZED:
                self._initialized.set()
            return

        if event.event_type in self.CONNECTIVE_EVENTS:
            event_name = self.CONNECTIVE_EVENTS[event.event_type]
            self._log.info('Connective event %s happened', event_name)
            return

        if event.event_type in (
                TreeEvent.NODE_ADDED,
                TreeEvent.NODE_UPDATED,
                TreeEvent.NODE_REMOVED):
            # NOTE We are able to do something on the tree cache now
            return

    def _handle_fault(self, error):
        self._log.exception(error)


def example():
    client = KazooClient()
    client.start()
    tree_cache = TreeCache(client, '/kazoo-treecache-example')
    tree_listener = TreeCacheListener(tree_cache)
    tree_listener.start()  # block until the tree cache is initialized
    path, data, stat = tree_cache.get_data('/kazoo-treecache-example')
    print('\n=> {0} = {1!r}; {2}'.format(path, data, stat))


if __name__ == '__main__':
    logging.basicConfig(level=logging.DEBUG)
    example()