#       Licensed to the Apache Software Foundation (ASF) under one
#       or more contributor license agreements.  See the NOTICE file
#       distributed with this work for additional information
#       regarding copyright ownership.  The ASF licenses this file
#       to you under the Apache License, Version 2.0 (the
#       "License"); you may not use this file except in compliance
#       with the License.  You may obtain a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0
#
#       Unless required by applicable law or agreed to in writing,
#       software distributed under the License is distributed on an
#       "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
#       KIND, either express or implied.  See the License for the
#       specific language governing permissions and limitations
#       under the License.

import logging
import pymongo

from ming import Session
from ming.orm.base import state
from ming.orm.ormsession import ThreadLocalORMSession, SessionExtension
from contextlib import contextmanager

from allura.lib.utils import chunked_list
from allura.tasks import index_tasks

log = logging.getLogger(__name__)


class ArtifactSessionExtension(SessionExtension):

    def __init__(self, session):
        SessionExtension.__init__(self, session)
        self.objects_added = []
        self.objects_modified = []
        self.objects_deleted = []

    def before_flush(self, obj=None):
        if obj is None:
            self.objects_added = list(self.session.uow.new)
            self.objects_modified = list(self.session.uow.dirty)
            self.objects_deleted = list(self.session.uow.deleted)
        else:  # pragma no cover
            st = state(obj)
            if st.status == st.new:
                self.objects_added = [obj]
            elif st.status == st.dirty:
                self.objects_modified = [obj]
            elif st.status == st.deleted:
                self.objects_deleted = [obj]

    def after_flush(self, obj=None):
        "Update artifact references, and add/update this artifact to solr"
        if not getattr(self.session, 'disable_artifact_index', False):
            from pylons import app_globals as g
            from .index import ArtifactReference, Shortlink
            from .session import main_orm_session
            # Ensure artifact references & shortlinks exist for new objects
            arefs = []
            try:
                arefs = [
                    ArtifactReference.from_artifact(o)
                    for o in self.objects_added + self.objects_modified]
                for obj in self.objects_added + self.objects_modified:
                    Shortlink.from_artifact(obj)
                # Flush shortlinks
                main_orm_session.flush()
            except Exception:
                log.exception(
                    "Failed to update artifact references. Is this a borked project migration?")
            self.update_index(self.objects_deleted, arefs)
            for obj in self.objects_added:
                g.zarkov_event('create', extra=obj.index_id())
            for obj in self.objects_modified:
                g.zarkov_event('modify', extra=obj.index_id())
            for obj in self.objects_deleted:
                g.zarkov_event('delete', extra=obj.index_id())
        self.objects_added = []
        self.objects_modified = []
        self.objects_deleted = []

    def update_index(self, objects_deleted, arefs):
        # Post delete and add indexing operations
        if objects_deleted:
            index_tasks.del_artifacts.post(
                [obj.index_id() for obj in objects_deleted])
        add_task = None
        if arefs:
            add_task = index_tasks.add_artifacts.post([aref._id for aref in arefs])
        try:
            l = logging.getLogger('allura.debug7047')
            from tg import request
            task = request.environ.get('task')
            if task and task.task_name == 'forgetracker.tasks.bulk_edit':
                l.debug('session: %s %s', self.session.impl.db, self.session)
                l.debug('arefs: %s', arefs)
                l.debug('objects_added: %s', [o._id for o in self.objects_added])
                l.debug('objects_modified: %s', [o._id for o in self.objects_modified])
                l.debug('objects_deleted: %s', [o._id for o in self.objects_deleted])
                l.debug('add_artifacts task: %s', add_task)
        except:
            pass


class BatchIndexer(ArtifactSessionExtension):

    """
    Tracks needed search index operations over the life of a
    :class:`ming.odm.session.ThreadLocalODMSession` session, and performs them
    in a batch when :meth:`flush` is called.
    """
    to_delete = set()
    to_add = set()

    def __init__(self, session):
        ArtifactSessionExtension.__init__(self, session)

    def update_index(self, objects_deleted, arefs_added):
        """
        Caches adds and deletes for handling later.  Called after each flush of
        the parent session.

        :param objects_deleted: :class:`allura.model.artifact.Artifact`
          instances that were deleted in the flush.

        :param arefs_added: :class:`allura.model.artifact.ArtifactReference`
          instances for all ``Artifact`` instances that were added or modified
          in the flush.
        """

        from .index import ArtifactReference
        del_index_ids = [obj.index_id() for obj in objects_deleted]
        deleted_aref_ids = [aref._id for aref in
                            ArtifactReference.query.find(dict(_id={'$in': del_index_ids}))]
        cls = self.__class__
        cls.to_add -= set(deleted_aref_ids)
        cls.to_delete |= set(del_index_ids)
        cls.to_add |= set([aref._id for aref in arefs_added])

    @classmethod
    def flush(cls):
        """
        Creates indexing tasks for cached adds and deletes, and resets the
        caches.

        .. warning:: This method is NOT called automatically when the parent
           session is flushed. It MUST be called explicitly.
        """
        # Post in chunks to avoid overflowing the max BSON document
        # size when the Monq task is created:
        # cls.to_delete - contains solr index ids which can easily be over
        #                 100 bytes. Here we allow for 160 bytes avg, plus
        #                 room for other document overhead.
        # cls.to_add - contains BSON ObjectIds, which are 12 bytes each, so
        #              we can easily put 1m in a doc with room left over.
        if cls.to_delete:
            for chunk in chunked_list(list(cls.to_delete), 100 * 1000):
                cls._post(index_tasks.del_artifacts, chunk)

        if cls.to_add:
            for chunk in chunked_list(list(cls.to_add), 1000 * 1000):
                cls._post(index_tasks.add_artifacts, chunk)
        cls.to_delete = set()
        cls.to_add = set()

    @classmethod
    def _post(cls, task_func, chunk):
        """
        Post task, recursively splitting and re-posting if the resulting
        mongo document is too large.
        """
        try:
            task_func.post(chunk)
        except pymongo.errors.InvalidDocument as e:
            # there are many types of InvalidDocument, only recurse if its
            # expected to help
            if str(e).startswith('BSON document too large'):
                cls._post(task_func, chunk[:len(chunk) // 2])
                cls._post(task_func, chunk[len(chunk) // 2:])
            else:
                raise


@contextmanager
def substitute_extensions(session, extensions=None):
    """
    Temporarily replace the extensions on a
    :class:`ming.odm.session.ThreadLocalODMSession` session.
    """
    original_exts = session._kwargs.get('extensions', [])

    def _set_exts(exts):
        session.flush()
        session.close()
        session._kwargs['extensions'] = exts
    _set_exts(extensions or [])
    yield session
    _set_exts(original_exts)


main_doc_session = Session.by_name('main')
project_doc_session = Session.by_name('project')
task_doc_session = Session.by_name('task')
main_orm_session = ThreadLocalORMSession(main_doc_session)
project_orm_session = ThreadLocalORMSession(project_doc_session)
task_orm_session = ThreadLocalORMSession(task_doc_session)
artifact_orm_session = ThreadLocalORMSession(
    doc_session=project_doc_session,
    extensions=[ArtifactSessionExtension])
repository_orm_session = ThreadLocalORMSession(
    doc_session=main_doc_session,
    extensions=[])