# coding=utf8

"""
    flask-idempotent2
    ~~~~~~~~~~~~~~~~~

    Redis based automatic idempotent support for sqlalchemy based flask
    applications.

    :copyright: (c) 2016 by Chao Wang (hit9).
    :license: BSD, see LICENSE for more details.
"""

import functools
import hashlib
import pickle

from flask import g, request, session as flask_session, Response, \
    _app_ctx_stack
from sqlalchemy import event
from sqlalchemy.inspection import inspect
from redis.lock import Lock


__version__ = '0.0.6'


def gen_keyfunc(endpoint=True, path=True, method=True, query_string=True,
                data=True, headers=None, session=True, content_type=True,
                content_length=True, remote_addr=True, use_checksum=True):
    """Generates a `keyfunc` that distinguishes requests on different
    dimensions. Requests is considered to be different if on the dimension if
    it is set.

    :param use_checksum: Defaults to ``True``. When ``True``, use checksum
       string instead of original key.

    """

    def keyfunc():
        dimensions = {}
        if endpoint:
            dimensions['endpoint'] = request.endpoint
        if path:
            dimensions['path'] = request.path
        if method:
            dimensions['method'] = request.method
        if query_string:
            dimensions['query_string'] = request.query_string
        if data:
            dimensions['data'] = request.data
        if headers:
            d = {}
            for name in headers:
                d[name] = request.headers.get(name, None)
            dimensions['headers'] = str(sorted(d.items()))
        if session:
            dimensions['session'] = str(sorted(flask_session.items()))
        if content_type:
            dimensions['content_type'] = request.content_type
        if content_length:
            dimensions['content_length'] = request.content_length
        if remote_addr:
            dimensions['remote_addr'] = request.headers.get(
                'X-Forwarded-For',
                request.remote_addr)
        origin_key = str(sorted(dimensions.items()))
        if use_checksum:
            # Use hashed stringify dimensions
            sha = hashlib.sha1()
            sha.update(origin_key.encode('utf8'))
            return sha.hexdigest()
        return origin_key
    return keyfunc


class Idempotent(object):
    """The idempotent object implements idempotent support for sqlalchemy based
    flask applications. It acts as a central registry for idempotent view
    functions. Example::

        from flask import Flask
        import redis
        from .models import Session

        app = Flask(__name__)
        redis_client = redis.StrictRedis()
        idempotent = Idempotent(app, redis_client, Session)

        @app.route('/resource', methods=['PUT'])
        @idempotent.cache(5)
        def create_or_update_resource():
            pass

        @app.route('/another_resource', methods=['PUT'])
        @idempotent.lock(5)
        def create_or_update_another_resource():
            pass

    :param app: A flask application object to work with.
    :param redis: A redis client constructed by `redis-py`.
    :param session_factory: A sqlalchemy session factory.
    :param default_keyfunc: `keyfunc` for all view functions by default.
       Default to ``gen_keyfunc()``.
    :param default_redis_lock_class: The default redis-py lock class to use.
       Defaults to be ``redis.lock.Lock``.
    :param redis_key_prefix: The key prefix string to use in redis. Default to
       ``'idempotent'``.
    :param g_: The global context to use, defaults to be ``None``. Setting to
       ``None`` means ``flask.g``.
    """

    REDIS_KEY_PREFIX_CACHE = 'cache'
    REDIS_KEY_PREFIX_LOCK = 'lock'

    def __init__(self, app, redis, session_factory, default_keyfunc=None,
                 redis_key_prefix='idempotent', default_redis_lock_class=None,
                 g_=None):
        self.app = app
        self.redis = redis
        self.session_factory = session_factory
        self.default_keyfunc = default_keyfunc or gen_keyfunc()
        self.redis_key_prefix = redis_key_prefix
        if default_redis_lock_class is None:
            default_redis_lock_class = Lock
        self.default_redis_lock_class = default_redis_lock_class
        self.g_ = g_ or g

        # This flag would be updated on first ``idempotent.cache`` registered.
        self.sqlalchemy_events_registered = False

    ###
    # Public API
    ###

    def cache(self, timeout, keyfunc=None):
        """Produce a decorator to wrap view function to cache responses for
        given seconds. Example::

            @app.route('/api', methods=['PUT'])
            @idempotent.cache(5)
            def api():
                pass

        :param timeout: Cache expiration in redis, in seconds.
        :param keyfunc: `keyfunc` to distinguish requests. Which can be
           generated by `gen_keyfunc`. If passed, it will be used instead of
           `default_keyfunc`.

        Notes that if a function won't be wrapped again if it's already wrapped
        by the ``idempotent.cache``.

        Raises ``TypeError`` if this decorator tries to apply on a ``lock``
        decorated function. The reason is that the idempotent makes the lock
        miss. Instead, the ``lock`` decorator should be used outside the
        ``cache``.
        """

        def decorator(func):
            if getattr(func, '_idempotent_lock_wrapped', False):
                raise TypeError(
                    "Can't decorate `lock` decorated function by `cache`,"
                    " in turn please!")
            return self.wrap_cache(func, timeout, keyfunc=keyfunc)
        return decorator

    def lock(self, timeout, lock_class=None, blocking_timeout=None,
             keyfunc=None):
        """Produce a decorator to wrap view function to lock concurrency
        requests for given seconds. Example::

            @app.route('/api', methods=['PUT'])
            @idempotent.lock(5)
            def api():
                pass

        :param timeout: Lock expiration in redis, in seconds.
        :param lock_class: Redis-py lock class to use. If passed, it will be
           used instead of `default_redis_lock_class`.
        :param blocking_timeout: Max blocking or wait time in seconds. Defaults
           to ``None``, which indicates continue to try forever.
        :param keyfunc: `keyfunc` to distinguish requests. Which can be
           generated by `gen_keyfunc`. If passed, it will be used instead of
           `default_keyfunc`.

        Notes that if a function won't be wrapped again if it's already wrapped
        by the ``idempotent.lock``.
        """

        def decorator(func):
            return self.wrap_lock(func, timeout, lock_class=lock_class,
                                  blocking_timeout=blocking_timeout,
                                  keyfunc=keyfunc)
        return decorator

    ###
    # Utils
    ###

    def format_redis_key(self, prefix, origin_key):
        """Format `origin_key` with redis prefix and app name.

        :param origin_key: The original key without prefix.
        :param prefix: The prefix string either ``REDIS_KEY_PREFIX_CACHE`` or
           ``REDIS_KEY_PREFIX_LOCK``.
        """
        return '{0}:{1}:{2}:{3}'.format(self.redis_key_prefix, self.app.name,
                                        prefix, origin_key)

    ###
    # SQLAlchemy Events (Cache)
    ###

    def register_sqlalchemy_events(self):
        """Register sqlalchemy event listeners.

        Note that registered hook functions will be called only if current
        invoked view function is marked as idempotent (by `cache`).
        """
        event.listen(self.session_factory, 'before_flush',
                     self.record_changed_instances)
        event.listen(self.session_factory, 'before_commit',
                     self.record_changed_instances)
        event.listen(self.session_factory, 'after_rollback',
                     self.clear_changed_instances)
        event.listen(self.session_factory, 'after_commit',
                     self.record_committed_changes)

    def should_call_sqlalchemy_hook(self):
        """Returns ``True`` if current invoked view function is marked as
        idempotent (by `cache`).
        Returns ``False`` if current progress is out side flask request
        context.
        """
        if _app_ctx_stack.top is None:  # Outside flask
            return False
        view_function = self.app.view_functions.get(request.endpoint, None)
        return view_function and getattr(view_function,
                                         '_idempotent_cache_wrapped', False)

    def record_changed_instances(self, session, flush_context=None,
                                 instances=None):
        """Record changed resource instances to ``self.g_`` before sqlalchemy
        session flush or commit. The ``self.g_`` is request scoped, it's new
        for each request.

        Note that the ``after_commit()`` hook is not per-flush, that is, the
        Session can emit SQL to the database many times within the scope of a
        transaction.
        """
        if not self.should_call_sqlalchemy_hook():
            return
        if not hasattr(self.g_, '_idempotent_changes'):
            self.g_._idempotent_changes = []
        # New instances.
        for instance in session.new:
            self.g_._idempotent_changes.append(instance)
        # Dirty instances.
        for instance in session.dirty:
            self.g_._idempotent_changes.append(instance)
        # Deleted instances.
        for instance in session.deleted:
            self.g_._idempotent_changes.append(instance)

    def clear_changed_instances(self, session):
        """Clear changed instances after sqlalchemy session rollback.
        """
        if not self.should_call_sqlalchemy_hook():
            return
        self.g_._idempotent_changes = []
        self.g_._idempotent_committed_changes = []

    def record_committed_changes(self, session):
        """Record committed changes to ``self.g_`` after sqlalchemy session
        commit.
        """
        if not self.should_call_sqlalchemy_hook():
            return
        changes = getattr(self.g_, '_idempotent_changes', None)
        if changes:
            self.g_._idempotent_committed_changes = changes
            self.g_._idempotent_changes = []

    ###
    # Idempotent Cache
    ###

    def wrap_cache(self, func, timeout, keyfunc=None):
        if getattr(func, '_idempotent_cache_wrapped', False):
            # Register only once.
            return

        if self.sqlalchemy_events_registered is False:
            # Lazy register sqlalchemy events for cache.
            self.register_sqlalchemy_events()
            self.sqlalchemy_events_registered = True

        if keyfunc is None:
            keyfunc = self.default_keyfunc

        @functools.wraps(func)
        def _wrapped(*args, **kwargs):
            idempotent_id = keyfunc()
            cached_response = self.get_cached_response(idempotent_id)
            if cached_response is not None:
                return cached_response
            response = func(*args, **kwargs)
            self.set_response_cache(idempotent_id, response, timeout)
            return response
        _wrapped._idempotent_cache_wrapped = True
        return _wrapped

    def format_cache_redis_key(self, origin_key):
        return self.format_redis_key(self.REDIS_KEY_PREFIX_CACHE, origin_key)

    def loads_cached_value(self, value):
        """Loads cached value into flask response and list of resource
        instances. Returns a tuple in format of `(response, instances)`.

        :param value: String value read from redis.
        """
        return pickle.loads(value)

    def dumps_response_and_changed_instances(self, response):
        """Dumps response and changed instances into string to used as redis
        cached value. The serialization is based on ``pickle``. Returns a
        tuple in format of `(picked, instance_strings)`.

        :param response: The `flask.Response` instance to be dump as string.
        """
        data = []

        # Given `response` may be a ``string``, ``tuple`` or ``Response``.
        if isinstance(response, Response):
            # Freeze response instances before ``pickle``.
            response.freeze()
        data.append(response)

        instances = set()

        # Committed instances, in format of ``[string, ...]``. Each
        # ``instance_str`` is constructed by resource name and primary key
        # value, e.g. ``'User:1'``.
        for instance in getattr(self.g_, '_idempotent_committed_changes', []):
            instance_state = inspect(instance)
            if instance_state.persistent:
                cls = instance_state.class_
                name = cls.__name__  # resource name
                pk_column = inspect(cls).primary_key[0]  # pk name
                pk = getattr(instance, pk_column.name, None)
                instance_str = '{0}:{1}'.format(name, pk)
            instances.add(instance_str)

        data.append(instances)
        return pickle.dumps(data), instances

    def get_cached_response(self, idempotent_id):
        """Get idempotent request cache before view function is called. Returns
        ``None`` on cache miss.

        :param idempotent_id: Idempotent reuqest string id produced by
           `keyfunc`.
        """
        # Get cached response and affected instances.
        key = self.format_cache_redis_key(idempotent_id)
        value = self.redis.get(key)
        if not value:
            return  # Miss

        # Get affected instances and validate
        response, instances = self.loads_cached_value(value)
        if instances:
            keys = [self.format_cache_redis_key(k) for k in instances]
            values = self.redis.mget(keys)
            for value in values:
                if value.decode('utf8') != idempotent_id:
                    self.redis.delete(key)
                    return  # Miss
        return response  # Hit

    def set_response_cache(self, idempotent_id, response, timeout):
        """Set idempotent request cache after view function is actually called.

        :param response: The response to cache.
        :param timeout: Cache expiration in seconds.
        :param idempotent_id: Idempotent reuqest string id produced by
           `keyfunc`.
        """
        value, instances = self.dumps_response_and_changed_instances(response)

        # Cache response by `idempotent_id`.
        key = self.format_cache_redis_key(idempotent_id)
        self.redis.setex(key, timeout, value)

        # Cache affected instances.
        for instance_str in instances:
            self.redis.setex(self.format_cache_redis_key(instance_str),
                             timeout, idempotent_id)

    ###
    # Idempotent Lock
    ###

    def format_lock_redis_key(self, origin_key):
        return self.format_redis_key(self.REDIS_KEY_PREFIX_LOCK, origin_key)

    def wrap_lock(self, func, timeout, lock_class=None, blocking_timeout=None,
                  keyfunc=None):
        if getattr(func, '_idempotent_lock_wrapped', False):
            # Register only once.
            return

        if keyfunc is None:
            keyfunc = self.default_keyfunc
        if lock_class is None:
            lock_class = self.default_redis_lock_class

        @functools.wraps(func)
        def _wrapped(*args, **kwargs):
            key = self.format_lock_redis_key(keyfunc())
            with self.redis.lock(key, timeout=timeout, lock_class=lock_class,
                                 blocking_timeout=blocking_timeout):
                return func(*args, **kwargs)
        _wrapped._idempotent_lock_wrapped = True
        return _wrapped
