#
# Copyright (c) 2007 Canonical
# Copyright (c) 2007 Thomas Herve <thomas@nimail.org>
#
# This file is part of Storm Object Relational Mapper.
#
# Storm is free software; you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as
# published by the Free Software Foundation; either version 2.1 of
# the License, or (at your option) any later version.
#
# Storm is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#

"""
Store wrapper and custom thread runner.
"""

from threading import Thread
from Queue import Queue

from storm.store import Store, AutoReload
from storm.info import get_obj_info
from elisa.extern.twisted_storm.wrapper import partial, DeferredResult, DeferredResultSet

from elisa.core.utils.defer import Deferred, deferredGenerator
from elisa.core.utils.defer import waitForDeferred, succeed, fail
from twisted.python.failure import Failure

from storm.info import get_cls_info, get_obj_info


class AlreadyStopped(Exception):
    """
    Except raised when a store is stopped multiple time.
    """



class StoreThread(Thread):
    """
    A thread class that wraps methods calls and fires deferred in the reactor
    thread.
    """
    STOP = object()

    def __init__(self):
        """
        Initialize the thread, and create a L{Queue} to stack jobs.
        """
        Thread.__init__(self)
        self.setDaemon(True)
        self._queue = Queue()
        self._stop_deferred = None
        self.stopped = False


    def defer_to_thread(self, f, *args, **kwargs):
        """
        Run the given function in the thread, wrapping the result with a
        L{Deferred}.

        @return: a deferred whose result will be the result of the function.
        @rtype: C{Deferred}
        """
        if self.stopped:
            # to prevent having pending calls after the thread got stopped
            return fail(AlreadyStopped(f))
        d = Deferred()
        self._queue.put((d, f, args, kwargs))
        return d


    def run(self):
        """
        Main execution loop: retrieve jobs from the queue and run them.
        """
        from twisted.internet import reactor
        o = self._queue.get()
        while o is not self.STOP:
            d, f, args, kwargs = o
            if d.called:
                o = self._queue.get()
                continue
            try:
                result = f(*args, **kwargs)
            except:
                f = Failure()
                reactor.callFromThread(d.errback, f)
            else:
                reactor.callFromThread(d.callback, result)
            o = self._queue.get()
        reactor.callFromThread(self._stop_deferred.callback, None)


    def stop(self):
        """
        Stop the thread.
        """
        if self.stopped:
            return self._stop_deferred
        self._stop_deferred = Deferred()
        self._queue.put(self.STOP)
        self.stopped = True
        return self._stop_deferred



class DeferredStore(object):
    """
    A wrapper around L{Store} to have async operations.
    """
    store = None

    def __init__(self, database, auto_reload=True):
        """
        @param database: instance of database providing connection, used to
            instantiate the store later.
        @type database: L{storm.database.Database}
        @param auto_reload: should the store reload all the lazy
                            attributes kept in the store cache. This is
                            a performance tweak option that you should
                            only set to False if you don't want to do any
                            raw sql inside a transaction. If you set it
                            to False but do raw sql request be aware of
                            the fact that the cached object might not be
                            in a state reflecting the database.
        @type auto_reload: bool
        """
        self.thread = StoreThread()
        self.database = database
        self.started = False
        self.auto_reload = auto_reload


    def start(self):
        """
        Start the store.

        @return: a deferred that will fire once the store is started.
        """
        if not self.started:
            self.started = True
            self.thread.start()
            # Add a event trigger to be sure that the thread is stopped
            from twisted.internet import reactor
            reactor.addSystemEventTrigger(
                "before", "shutdown", self.stop)
            return self.thread.defer_to_thread(Store, self.database
                ).addCallback(self._got_store)
        else:
            raise RuntimeError("Already started")


    def _got_store(self, store):
        """
        Internal method called when the store is created, initializing most of
        the API methods.
        """
        self.store = store
        self._choose_decorate_method()
        # Maybe not ?
        self.store._deferredStore = self

        for methodName in ("add", "remove", "reload", "flush", "rollback"):
            method = partial(self.thread.defer_to_thread,
                             getattr(self.store, methodName))
            setattr(self, methodName, method)

        # disable lazy resolving
        self.store.invalidate = self.invalidate
        self.store._enable_lazy_resolving = lambda x: None

        # set up some methods we want to have auto reload on
        for methodName in ("flush", "add", "rollback"):
            method = self.decorate_method(getattr(self.store, methodName))
            setattr(self.store, methodName, method)


    def _choose_decorate_method(self):
        """
        compability layer that helps to choose if the much more enhanced trunk
        method should be used as decorate_method or the one for the last release
        """
        if hasattr(self.store, "_implicit_flush_block_count"):
            # we have a trunk
            self.decorate_method = self.decorate_method_trunk
        else:
            self.decotare_method = self.decorate_method_release

    def decorate_method_trunk(self, method):
         def decorate(*args,**kw):
            result = method(*args, **kw)
            saved = self.store._implicit_flush_block_count
            self.store._implicit_flush_block_count = 1
            self._auto_reload_alives()
            self.store._implicit_flush_block_count = saved
            return result
         return decorate

    def decorate_method_release(self, method):
        def decorate(*args,**kw):
            result = method(*args, **kw)
            self._auto_reload_alives()
            return result
        return decorate

    def get(self, cls, key):
        def _get():
            result = self.store.get(cls, key)
            if result is not None:
                obj_info = get_obj_info(result)
                self.store._resolve_lazy_value(obj_info, None, AutoReload)
            return result
        return self.thread.defer_to_thread(_get)

    def _auto_reload_alives(self):
        if self.auto_reload:
            self._reload_alives()

    def _reload_alives(self):
        for obj_info in self.store._alive.itervalues():
            self.store._resolve_lazy_value(obj_info, None, AutoReload)

    def commit(self):
        def _commit():
            result = self.store.commit()
            self._auto_reload_alives()
            return result

        return self.thread.defer_to_thread(self.store.commit)

    def execute(self, *args, **kwargs):
        """
        Wrapper around C{execute} to have a C{DeferredResult} instead of the
        standard L{storm.database.Result} object.
        """
        if self.store is None:
            raise RuntimeError("Store not started")
        return self.thread.defer_to_thread(
            self.store.execute, *args, **kwargs
            ).addCallback(self._cb_execute)


    def _cb_execute(self, result):
        """
        Wrap the result with a C{DeferredResult}.
        """
        return DeferredResult(self.thread, result)


    def find(self, *args, **kwargs):
        """
        Wrapper around C{find}.
        """
        if self.store is None:
            raise RuntimeError("Store not started")
        return self.thread.defer_to_thread(
            self.store.find, *args, **kwargs
            ).addCallback(self._cb_find)


    def _cb_find(self, resultSet):
        """
        Wrap the result set with a C{DeferredResultSet}.
        """
        return DeferredResultSet(self.thread, resultSet)


    def stop(self):
        """
        Stop the store.
        """
        return self.thread.stop()


    def invalidate(self, obj=None):
        """
        Invalidate isn't meaningful on the store for now.
        """



class StorePool(object):
    """
    A pool of started stores, maintaining persistent connections.
    """
    started = False
    store_factory = DeferredStore

    def __init__(self, database, min_stores=0, max_stores=10):
        """
        @param database: instance of database providing connection, used to
            instantiate the store later.
        @type database: L{storm.database.Database}

        @param min_stores: initial number of stores.
        @type min_stores: C{int}

        @param max_stores: maximum number of stores.
        @type max_stores: C{int}
        """
        self.database = database
        self.min_stores = min_stores
        self.max_stores = max_stores
        self._stores = []
        self._stores_created = 0
        self._pending_get = []


    def start(self):
        """
        Start the pool.
        """
        if self.started:
            raise RuntimeError("Already started")
        self.started = True
        return self.adjust_size()


    def stop(self):
        """
        Stop the pool: this is not a total stop, it just try to kill the
        current available stores.
        """
        return self.adjust_size(0, 0)


    def start_store(self):
        """
        Create a start a new store.
        """
        store = self.store_factory(self.database)
        return store.start().addCallback(self._cbStartAStore, store)


    def _cbStartAStore(self, ign, store):
        """
        Add the created store to the list of available stores.
        """
        self._stores_created += 1
        self._stores.append(store)


    def stop_store(self):
        """
        Stop a store and remove it from the available stores.
        """
        self._stores_created -= 1
        store = self._stores.pop()
        return store.stop()


    def adjust_size(self, min_stores=None, max_stores=None):
        """
        Change the number of available stores, shrinking or raising as
        necessary.
        """
        if min_stores is None:
            min_stores = self.min_stores
        if max_stores is None:
            max_stores = self.max_stores

        if min_stores < 0:
            raise ValueError('minimum is negative')
        if min_stores > max_stores:
            raise ValueError('minimum is greater than maximum')

        self.min_stores = min_stores
        self.max_stores = max_stores
        if not self.started:
            return

        # Kill of some stores if we have too many.
        while self._stores_created > self.max_stores and self._stores:
            wfd = waitForDeferred(self.stop_store())
            yield wfd
            wfd.getResult()
        # Start some stores if we have too few.
        while self._stores_created < self.min_stores:
            wfd = waitForDeferred(self.start_store())
            yield wfd
            wfd.getResult()

    adjust_size = deferredGenerator(adjust_size)


    def get(self):
        """
        Return a started store from the pool, or start a new one if necessary.
        A store retrieve by this way should be put back using the put
        method, or it won't be used anymore.
        """
        if not self.started:
            raise RuntimeError("Not started")
        if self._stores:
            store = self._stores.pop()
            return succeed(store)
        elif self._stores_created < self.max_stores:
            return self.start_store().addCallback(self._cb_get)
        else:
            # Maybe all stores are consumed?
            return self.adjust_size().addCallback(self._cb_get)


    def _cb_get(self, ign):
        """
        If the previous operation added a store, return it, or return a pending
        C{Deferred}.
        """
        if self._stores:
            store = self._stores.pop()
            return store
        else:
            # All stores are in used, wait
            d = Deferred()
            self._pending_get.append(d)
            return d


    def put(self, store):
        """
        Make a store available again.

        This should be done explicitely to have the store back in the pool.
        The good way to use the pool is this:

        >>> d1 = pool.get()

        >>> # d1 callback with a store
        >>> d2 = store.add(foo)
        >>> d2.addCallback(doSomething).addErrback(manageErrors)
        >>> d2.addBoth(lambda x: pool.put(store))
        """
        return store.rollback().addCallback(self._cb_put, store)


    def _cb_put(self, ign, store):
        """
        Once the rollback has finished, the store is really available.
        """
        if self._pending_get:
            # People are waiting, fire with the store
            d = self._pending_get.pop(0)
            d.callback(store)
        else:
            self._stores.append(store)
