# -*- coding: utf-8 -*-
'''
Redis
=====

Redis plugin for the Salt caching subsystem.

.. versionadded:: 2017.7.0

As Redis provides a simple mechanism for very fast key-value store, in order to
privde the necessary features for the Salt caching subsystem, the following
conventions are used:

- A Redis key consists of the bank name and the cache key separated by ``/``, e.g.:
  ``$KEY_minions/alpha/stuff`` where ``minions/alpha`` is the bank name
  and ``stuff`` is the key name.
- As the caching subsystem is organised as a tree, we need to store the caching
  path and identify the bank and its offspring.  At the same time, Redis is
  linear and we need to avoid doing ``keys <pattern>`` which is very
  inefficient as it goes through all the keys on the remote Redis server.
  Instead, each bank hierarchy has a Redis SET associated which stores the list
  of sub-banks. By default, these keys begin with ``$BANK_``.
- In addition, each key name is stored in a separate SET of all the keys within
  a bank. By default, these SETs begin with ``$BANKEYS_``.

For example, to store the key ``my-key`` under the bank ``root-bank/sub-bank/leaf-bank``,
the following hierarchy will be built:

.. code-block:: text

    127.0.0.1:6379> SMEMBERS $BANK_root-bank
    1) "sub-bank"
    127.0.0.1:6379> SMEMBERS $BANK_root-bank/sub-bank
    1) "leaf-bank"
    127.0.0.1:6379> SMEMBERS $BANKEYS_root-bank/sub-bank/leaf-bank
    1) "my-key"
    127.0.0.1:6379> GET $KEY_root-bank/sub-bank/leaf-bank/my-key
    "my-value"

There are three types of keys stored:

- ``$BANK_*`` is a Redis SET containing the list of banks under the current bank
- ``$BANKEYS_*`` is a Redis SET containing the list of keys under the current bank
- ``$KEY_*`` keeps the value of the key

These prefixes and the separator can be adjusted using the configuration options:

bank_prefix: ``$BANK``
    The prefix used for the name of the Redis key storing the list of sub-banks.

bank_keys_prefix: ``$BANKEYS``
    The prefix used for the name of the Redis keyt storing the list of keys under a certain bank.

key_prefix: ``$KEY``
    The prefix of the Redis keys having the value of the keys to be cached under
    a certain bank.

separator: ``_``
    The separator between the prefix and the key body.

The connection details can be specified using:

host: ``localhost``
    The hostname of the Redis server.

port: ``6379``
    The Redis server port.

cluster_mode: ``False``
    Whether cluster_mode is enabled or not

cluster.startup_nodes:
    A list of host, port dictionaries pointing to cluster members. At least one is required
    but multiple nodes are better

    .. code-block:: yaml

        cache.redis.cluster.startup_nodes
          - host: redis-member-1
            port: 6379
          - host: redis-member-2
            port: 6379

cluster.skip_full_coverage_check: ``False``
    Some cluster providers restrict certain redis commands such as CONFIG for enhanced security.
    Set this option to true to skip checks that required advanced privileges.

    .. note::

        Most cloud hosted redis clusters will require this to be set to ``True``

db: ``'0'``
    The database index.

    .. note::
        The database index must be specified as string not as integer value!

password:
    Redis connection password.

unix_socket_path:

    .. versionadded:: 2018.3.1

    Path to a UNIX socket for access. Overrides `host` / `port`.

Configuration Example:

.. code-block:: yaml

    cache.redis.host: localhost
    cache.redis.port: 6379
    cache.redis.db: '0'
    cache.redis.password: my pass
    cache.redis.bank_prefix: #BANK
    cache.redis.bank_keys_prefix: #BANKEYS
    cache.redis.key_prefix: #KEY
    cache.redis.separator: '@'

Cluster Configuration Example:

.. code-block:: yaml

    cache.redis.cluster_mode: true
    cache.redis.cluster.skip_full_coverage_check: true
    cache.redis.cluster.startup_nodes:
      - host: redis-member-1
        port: 6379
      - host: redis-member-2
        port: 6379
    cache.redis.db: '0'
    cache.redis.password: my pass
    cache.redis.bank_prefix: #BANK
    cache.redis.bank_keys_prefix: #BANKEYS
    cache.redis.key_prefix: #KEY
    cache.redis.separator: '@'
'''

from __future__ import absolute_import, print_function, unicode_literals

# Import stdlib
import logging

# Import third party libs
try:
    import redis
    from redis.exceptions import ConnectionError as RedisConnectionError
    from redis.exceptions import ResponseError as RedisResponseError
    HAS_REDIS = True
except ImportError:
    HAS_REDIS = False

try:
    from rediscluster import StrictRedisCluster
    HAS_REDIS_CLUSTER = True
except ImportError:
    HAS_REDIS_CLUSTER = False

# Import salt
from salt.ext.six.moves import range
from salt.exceptions import SaltCacheError

# -----------------------------------------------------------------------------
# module properties
# -----------------------------------------------------------------------------

__virtualname__ = 'redis'
__func_alias__ = {'list_': 'list'}

log = logging.getLogger(__file__)

_BANK_PREFIX = '$BANK'
_KEY_PREFIX = '$KEY'
_BANK_KEYS_PREFIX = '$BANKEYS'
_SEPARATOR = '_'

REDIS_SERVER = None

# -----------------------------------------------------------------------------
# property functions
# -----------------------------------------------------------------------------


def __virtual__():
    '''
    The redis library must be installed for this module to work.

    The redis redis cluster library must be installed if cluster_mode is True
    '''
    if not HAS_REDIS:
        return (False, "Please install the python-redis package.")
    if not HAS_REDIS_CLUSTER and _get_redis_cache_opts()['cluster_mode']:
        return (False, "Please install the redis-py-cluster package.")
    return __virtualname__


# -----------------------------------------------------------------------------
# helper functions -- will not be exported
# -----------------------------------------------------------------------------

def init_kwargs(kwargs):
    return {}


def _get_redis_cache_opts():
    '''
    Return the Redis server connection details from the __opts__.
    '''
    return {
        'host': __opts__.get('cache.redis.host', 'localhost'),
        'port': __opts__.get('cache.redis.port', 6379),
        'unix_socket_path': __opts__.get('cache.redis.unix_socket_path', None),
        'db': __opts__.get('cache.redis.db', '0'),
        'password': __opts__.get('cache.redis.password', ''),
        'cluster_mode': __opts__.get('cache.redis.cluster_mode', False),
        'startup_nodes': __opts__.get('cache.redis.cluster.startup_nodes', {}),
        'skip_full_coverage_check': __opts__.get('cache.redis.cluster.skip_full_coverage_check', False),
    }


def _get_redis_server(opts=None):
    '''
    Return the Redis server instance.
    Caching the object instance.
    '''
    global REDIS_SERVER
    if REDIS_SERVER:
        return REDIS_SERVER
    if not opts:
        opts = _get_redis_cache_opts()

    if opts['cluster_mode']:
        REDIS_SERVER = StrictRedisCluster(startup_nodes=opts['startup_nodes'],
                                          skip_full_coverage_check=opts['skip_full_coverage_check'])
    else:
        REDIS_SERVER = redis.StrictRedis(opts['host'],
                                   opts['port'],
                                   unix_socket_path=opts['unix_socket_path'],
                                   db=opts['db'],
                                   password=opts['password'])
    return REDIS_SERVER


def _get_redis_keys_opts():
    '''
    Build the key opts based on the user options.
    '''
    return {
        'bank_prefix': __opts__.get('cache.redis.bank_prefix', _BANK_PREFIX),
        'bank_keys_prefix': __opts__.get('cache.redis.bank_keys_prefix', _BANK_KEYS_PREFIX),
        'key_prefix': __opts__.get('cache.redis.key_prefix', _KEY_PREFIX),
        'separator': __opts__.get('cache.redis.separator', _SEPARATOR)
    }


def _get_bank_redis_key(bank):
    '''
    Return the Redis key for the bank given the name.
    '''
    opts = _get_redis_keys_opts()
    return '{prefix}{separator}{bank}'.format(
        prefix=opts['bank_prefix'],
        separator=opts['separator'],
        bank=bank
    )


def _get_key_redis_key(bank, key):
    '''
    Return the Redis key given the bank name and the key name.
    '''
    opts = _get_redis_keys_opts()
    return '{prefix}{separator}{bank}/{key}'.format(
        prefix=opts['key_prefix'],
        separator=opts['separator'],
        bank=bank,
        key=key
    )


def _get_bank_keys_redis_key(bank):
    '''
    Return the Redis key for the SET of keys under a certain bank, given the bank name.
    '''
    opts = _get_redis_keys_opts()
    return '{prefix}{separator}{bank}'.format(
        prefix=opts['bank_keys_prefix'],
        separator=opts['separator'],
        bank=bank
    )


def _build_bank_hier(bank, redis_pipe):
    '''
    Build the bank hierarchy from the root of the tree.
    If already exists, it won't rewrite.
    It's using the Redis pipeline,
    so there will be only one interaction with the remote server.
    '''
    bank_list = bank.split('/')
    parent_bank_path = bank_list[0]
    for bank_name in bank_list[1:]:
        prev_bank_redis_key = _get_bank_redis_key(parent_bank_path)
        redis_pipe.sadd(prev_bank_redis_key, bank_name)
        log.debug('Adding %s to %s', bank_name, prev_bank_redis_key)
        parent_bank_path = '{curr_path}/{bank_name}'.format(
            curr_path=parent_bank_path,
            bank_name=bank_name
        )  # this becomes the parent of the next
    return True


def _get_banks_to_remove(redis_server, bank, path=''):
    '''
    A simple tree tarversal algorithm that builds the list of banks to remove,
    starting from an arbitrary node in the tree.
    '''
    current_path = bank if not path else '{path}/{bank}'.format(path=path, bank=bank)
    bank_paths_to_remove = [current_path]
    # as you got here, you'll be removed

    bank_key = _get_bank_redis_key(current_path)
    child_banks = redis_server.smembers(bank_key)
    if not child_banks:
        return bank_paths_to_remove  # this bank does not have any child banks so we stop here
    for child_bank in child_banks:
        bank_paths_to_remove.extend(_get_banks_to_remove(redis_server, child_bank, path=current_path))
        # go one more level deeper
        # and also remove the children of this child bank (if any)
    return bank_paths_to_remove

# -----------------------------------------------------------------------------
# cache subsystem functions
# -----------------------------------------------------------------------------


def store(bank, key, data):
    '''
    Store the data in a Redis key.
    '''
    redis_server = _get_redis_server()
    redis_pipe = redis_server.pipeline()
    redis_key = _get_key_redis_key(bank, key)
    redis_bank_keys = _get_bank_keys_redis_key(bank)
    try:
        _build_bank_hier(bank, redis_pipe)
        value = __context__['serial'].dumps(data)
        redis_pipe.set(redis_key, value)
        log.debug('Setting the value for %s under %s (%s)', key, bank, redis_key)
        redis_pipe.sadd(redis_bank_keys, key)
        log.debug('Adding %s to %s', key, redis_bank_keys)
        redis_pipe.execute()
    except (RedisConnectionError, RedisResponseError) as rerr:
        mesg = 'Cannot set the Redis cache key {rkey}: {rerr}'.format(rkey=redis_key,
                                                                      rerr=rerr)
        log.error(mesg)
        raise SaltCacheError(mesg)


def fetch(bank, key):
    '''
    Fetch data from the Redis cache.
    '''
    redis_server = _get_redis_server()
    redis_key = _get_key_redis_key(bank, key)
    redis_value = None
    try:
        redis_value = redis_server.get(redis_key)
    except (RedisConnectionError, RedisResponseError) as rerr:
        mesg = 'Cannot fetch the Redis cache key {rkey}: {rerr}'.format(rkey=redis_key,
                                                                        rerr=rerr)
        log.error(mesg)
        raise SaltCacheError(mesg)
    if redis_value is None:
        return {}
    return __context__['serial'].loads(redis_value)


def flush(bank, key=None):
    '''
    Remove the key from the cache bank with all the key content. If no key is specified, remove
    the entire bank with all keys and sub-banks inside.
    This function is using the Redis pipelining for best performance.
    However, when removing a whole bank,
    in order to re-create the tree, there are a couple of requests made. In total:

    - one for node in the hierarchy sub-tree, starting from the bank node
    - one pipelined request to get the keys under all banks in the sub-tree
    - one pipeline request to remove the corresponding keys

    This is not quite optimal, as if we need to flush a bank having
    a very long list of sub-banks, the number of requests to build the sub-tree may grow quite big.

    An improvement for this would be loading a custom Lua script in the Redis instance of the user
    (using the ``register_script`` feature) and call it whenever we flush.
    This script would only need to build this sub-tree causing problems. It can be added later and the behaviour
    should not change as the user needs to explicitly allow Salt inject scripts in their Redis instance.
    '''
    redis_server = _get_redis_server()
    redis_pipe = redis_server.pipeline()
    if key is None:
        # will remove all bank keys
        bank_paths_to_remove = _get_banks_to_remove(redis_server, bank)
        # tree traversal to get all bank hierarchy
        for bank_to_remove in bank_paths_to_remove:
            bank_keys_redis_key = _get_bank_keys_redis_key(bank_to_remove)
            # Redis key of the SET that stores the bank keys
            redis_pipe.smembers(bank_keys_redis_key)  # fetch these keys
            log.debug(
                'Fetching the keys of the %s bank (%s)',
                bank_to_remove, bank_keys_redis_key
            )
        try:
            log.debug('Executing the pipe...')
            subtree_keys = redis_pipe.execute()  # here are the keys under these banks to be removed
            # this retunrs a list of sets, e.g.:
            # [set([]), set(['my-key']), set(['my-other-key', 'yet-another-key'])]
            # one set corresponding to a bank
        except (RedisConnectionError, RedisResponseError) as rerr:
            mesg = 'Cannot retrieve the keys under these cache banks: {rbanks}: {rerr}'.format(
                rbanks=', '.join(bank_paths_to_remove),
                rerr=rerr
            )
            log.error(mesg)
            raise SaltCacheError(mesg)
        total_banks = len(bank_paths_to_remove)
        # bank_paths_to_remove and subtree_keys have the same length (see above)
        for index in range(total_banks):
            bank_keys = subtree_keys[index]  # all the keys under this bank
            bank_path = bank_paths_to_remove[index]
            for key in bank_keys:
                redis_key = _get_key_redis_key(bank_path, key)
                redis_pipe.delete(redis_key)  # kill 'em all!
                log.debug(
                    'Removing the key %s under the %s bank (%s)',
                    key, bank_path, redis_key
                )
            bank_keys_redis_key = _get_bank_keys_redis_key(bank_path)
            redis_pipe.delete(bank_keys_redis_key)
            log.debug(
                'Removing the bank-keys key for the %s bank (%s)',
                bank_path, bank_keys_redis_key
            )
            # delete the Redis key where are stored
            # the list of keys under this bank
            bank_key = _get_bank_redis_key(bank_path)
            redis_pipe.delete(bank_key)
            log.debug('Removing the %s bank (%s)', bank_path, bank_key)
            # delete the bank key itself
    else:
        redis_key = _get_key_redis_key(bank, key)
        redis_pipe.delete(redis_key)  # delete the key cached
        log.debug(
            'Removing the key %s under the %s bank (%s)',
            key, bank, redis_key
        )
        bank_keys_redis_key = _get_bank_keys_redis_key(bank)
        redis_pipe.srem(bank_keys_redis_key, key)
        log.debug(
            'De-referencing the key %s from the bank-keys of the %s bank (%s)',
            key, bank, bank_keys_redis_key
        )
        # but also its reference from $BANKEYS list
    try:
        redis_pipe.execute()  # Fluuuush
    except (RedisConnectionError, RedisResponseError) as rerr:
        mesg = 'Cannot flush the Redis cache bank {rbank}: {rerr}'.format(rbank=bank,
                                                                          rerr=rerr)
        log.error(mesg)
        raise SaltCacheError(mesg)
    return True


def list_(bank):
    '''
    Lists entries stored in the specified bank.
    '''
    redis_server = _get_redis_server()
    bank_redis_key = _get_bank_redis_key(bank)
    try:
        banks = redis_server.smembers(bank_redis_key)
    except (RedisConnectionError, RedisResponseError) as rerr:
        mesg = 'Cannot list the Redis cache key {rkey}: {rerr}'.format(rkey=bank_redis_key,
                                                                       rerr=rerr)
        log.error(mesg)
        raise SaltCacheError(mesg)
    if not banks:
        return []
    return list(banks)


def contains(bank, key):
    '''
    Checks if the specified bank contains the specified key.
    '''
    redis_server = _get_redis_server()
    bank_redis_key = _get_bank_redis_key(bank)
    try:
        return redis_server.sismember(bank_redis_key, key)
    except (RedisConnectionError, RedisResponseError) as rerr:
        mesg = 'Cannot retrieve the Redis cache key {rkey}: {rerr}'.format(rkey=bank_redis_key,
                                                                           rerr=rerr)
        log.error(mesg)
        raise SaltCacheError(mesg)
