import asyncio
from motor.motor_asyncio import AsyncIOMotorDatabase, AsyncIOMotorCursor
from pymongo.errors import BulkWriteError, DuplicateKeyError

from ..builder import BaseBuilder
from ..document import DocumentImplementation
from ..embedded_document import EmbeddedDocumentImplementation
from ..data_proxy import missing
from ..data_objects import Reference
from ..exceptions import NotCreatedError, UpdateError, ValidationError, DeleteError
from ..fields import ReferenceField, ListField, EmbeddedField
from ..query_mapper import map_query

from .tools import cook_find_filter


class WrappedCursor(AsyncIOMotorCursor):

    __slots__ = ('raw_cursor', 'collection_cursor', 'document_cls')

    def __init__(self, document_cls, cursor, collection_cursor=False):
        # Such a cunning plan my lord !
        # We inherit from Cursor but don't call it __init__ because
        # we act as a proxy to the underlying raw_cursor
        WrappedCursor.raw_cursor.__set__(self, cursor)
        WrappedCursor.collection_cursor.__set__(self, collection_cursor)
        WrappedCursor.document_cls.__set__(self, document_cls)

    def __getattr__(self, name):
        return getattr(self.raw_cursor, name)

    def __setattr__(self, name, value):
        return setattr(self.raw_cursor, name, value)

    def clone(self):
        return WrappedCursor(self.document_cls, self.raw_cursor.clone(), self.collection_cursor)

    def next_object(self):
        raw = self.raw_cursor.next_object()
        return self.document_cls.build_from_mongo(raw, use_cls=True) if not self.collection_cursor else raw

    def each(self, callback):
        def wrapped_callback(result, error):
            if not error and result is not None and not self.collection_cursor:
                result = self.document_cls.build_from_mongo(result, use_cls=True)
            return callback(result, error)
        return self.raw_cursor.each(wrapped_callback)

    def to_list(self, length=None):
        raw_future = cooked_future = self.raw_cursor.to_list(length)
        if not self.collection_cursor:
            cooked_future = asyncio.Future()
            builder = self.document_cls.build_from_mongo

            def on_raw_done(fut):
                cooked_future.set_result([builder(e, use_cls=True) for e in fut.result()])

            raw_future.add_done_callback(on_raw_done)
        return cooked_future

    def count(self, **kwargs):
        cursor = self.delegate
        count_cond = {'filter': cursor._Cursor__spec}

        session = kwargs.get('session', cursor.session)
        if session:
            count_cond['session'] = session
        skip = kwargs.get('skip', cursor._Cursor__skip)
        if skip:
            count_cond['skip'] = skip
        limit = kwargs.get('limit', cursor._Cursor__limit)
        if limit:
            count_cond['limit'] = limit
        max_time_ms = kwargs.get('max_time_ms', cursor._Cursor__max_time_ms)
        if max_time_ms:
            count_cond['maxTimeMS'] = max_time_ms
        collation = kwargs.get('collation', cursor._Cursor__collation)
        if collation:
            count_cond['collation'] = collation
        hint = kwargs.get('hint', cursor._Cursor__hint)
        if hint:
            count_cond['hint'] = hint
        return self.collection.count_documents(**count_cond)


class MotorAsyncIODocument(DocumentImplementation):

    __slots__ = ()

    opts = DocumentImplementation.opts

    # Cook hooks into coroutines in order to allow them to return
    # either Future or regular return value.

    @asyncio.coroutine
    def __coroutined_pre_insert(self):
        return self.pre_insert()

    @asyncio.coroutine
    def __coroutined_pre_update(self):
        return self.pre_update()

    @asyncio.coroutine
    def __coroutined_pre_delete(self):
        return self.pre_delete()

    @asyncio.coroutine
    def __coroutined_post_insert(self, ret):
        return self.post_insert(ret)

    @asyncio.coroutine
    def __coroutined_post_update(self, ret):
        return self.post_update(ret)

    @asyncio.coroutine
    def __coroutined_post_delete(self, ret):
        return self.post_delete(ret)

    async def reload(self):
        """
        Retrieve and replace document's data by the ones in database.

        Raises :class:`umongo.exceptions.NotCreatedError` if the document
        doesn't exist in database.
        """
        if not self.is_created:
            raise NotCreatedError("Document doesn't exists in database")
        ret = await self.collection.find_one(self.pk)
        if ret is None:
            raise NotCreatedError("Document doesn't exists in database")
        self._data = self.DataProxy(load_data=False)
        self._data.from_mongo(ret, parent_instance=self)

    async def commit(self, io_validate_all=False, conditions=None):
        """
        Commit the document in database.
        If the document doesn't already exist it will be inserted, otherwise
        it will be updated.

        :param io_validate_all:
        :param conditions: only perform commit if matching record in db
            satisfies condition(s) (e.g. version number).
            Raises :class:`umongo.exceptions.UpdateError` if the
            conditions are not satisfied.
        :return: Update result dict returned by underlaying driver or
            ObjectId of the inserted document.
        """
        try:
            if self.is_created:
                ret = None
                if self.is_modified():
                    query = conditions or {}
                    query['_id'] = self.pk
                    # pre_update can provide additional query filter and/or
                    # modify the fields' values
                    additional_filter = await self.__coroutined_pre_update()
                    if additional_filter:
                        query.update(map_query(additional_filter, self.schema.fields))
                    self.required_validate()
                    await self.io_validate(validate_all=io_validate_all)
                    payload = self._data.to_mongo(update=True)
                    if payload:
                        ret = await self.collection.update_one(query, payload)
                        if ret.matched_count != 1:
                            raise UpdateError(ret)
                        await self.__coroutined_post_update(ret)

            elif conditions:
                raise RuntimeError('Document must already exist in database to use `conditions`.')
            else:
                await self.__coroutined_pre_insert()
                self.required_validate()
                await self.io_validate(validate_all=io_validate_all)
                payload = self._data.to_mongo(update=False)
                ret = await self.collection.insert_one(payload)
                # TODO: check ret ?
                self._data.set_by_mongo_name('_id', ret.inserted_id)
                self.is_created = True
                await self.__coroutined_post_insert(ret)
        except DuplicateKeyError as exc:
            # Need to dig into error message to find faulting index
            errmsg = exc.details['errmsg']
            for index in self.opts.indexes:
                if ('.$%s' % index.document['name'] in errmsg or
                        ' %s ' % index.document['name'] in errmsg):
                    keys = index.document['key'].keys()
                    if len(keys) == 1:
                        key = tuple(keys)[0]
                        msg = self.schema.fields[key].error_messages['unique']
                        raise ValidationError({key: msg})
                    else:
                        fields = self.schema.fields
                        # Compound index (sort value to make testing easier)
                        keys = sorted(keys)
                        raise ValidationError({k: fields[k].error_messages[
                            'unique_compound'].format(fields=keys) for k in keys})
            # Unknown index, cannot wrap the error so just reraise it
            raise
        self._data.clear_modified()
        return ret

    def delete(self, *args, **kwargs):
        """
        Alias of :meth:`remove` to enforce default api.
        """
        return self.remove(*args, **kwargs)

    async def remove(self, conditions=None):
        """
        Remove the document from database.

        :param conditions: Only perform delete if matching record in db
            satisfies condition(s) (e.g. version number).
            Raises :class:`umongo.exceptions.DeleteError` if the
            conditions are not satisfied.
        Raises :class:`umongo.exceptions.NotCreatedError` if the document
        is not created (i.e. ``doc.is_created`` is False)
        Raises :class:`umongo.exceptions.DeleteError` if the document
        doesn't exist in database.

        :return: Delete result dict returned by underlaying driver.
        """
        if not self.is_created:
            raise NotCreatedError("Document doesn't exists in database")
        query = conditions or {}
        query['_id'] = self.pk
        # pre_delete can provide additional query filter
        additional_filter = await self.__coroutined_pre_delete()
        if additional_filter:
            query.update(map_query(additional_filter, self.schema.fields))
        ret = await self.collection.delete_one(query)
        if ret.deleted_count != 1:
            raise DeleteError(ret)
        self.is_created = False
        await self.__coroutined_post_delete(ret)
        return ret

    @classmethod
    async def delete_many(cls, filter, *args, **kwargs):
        """
        Bulk remove the documents from database.
        """
        filter = cook_find_filter(cls, filter)
        ret = await cls.collection.delete_many(filter, *args, **kwargs)
        return ret

    @classmethod
    async def insert_many(cls, documents, io_validate_all=False, **kwargs):
        """
        Bulk insert the documents.
        """
        try:
            payloads = []
            async_validations = []
            for obj in documents:
                obj.required_validate()
                async_validations.append(obj.io_validate(validate_all=io_validate_all))
                payloads.append(obj._data.to_mongo(update=False))
            await asyncio.gather(*async_validations)

            ret = await cls.collection.insert_many(documents=payloads, **kwargs)
        except BulkWriteError as exc:
            raise ValidationError(
                "Bulk creating error. Inserted {count} document(s)".format(count=exc.details.get('nInserted', 0))
            )
        return ret

    @classmethod
    async def update_many(cls, filter, update, **kwargs):
        """
        Bulk update the documents.
        """
        filter = cook_find_filter(cls, filter)
        try:
            ret = await cls.collection.update_many(filter, update=update, **kwargs)
        except BulkWriteError as exc:
            raise ValidationError("Bulk updating error.")
        except DuplicateKeyError as exc:
            # Need to dig into error message to find faulting index
            errmsg = exc.details['errmsg']
            for index in cls.opts.indexes:
                if '.$%s' % index.document['name'] in errmsg or ' %s ' % index.document['name'] in errmsg:
                    keys = index.document['key'].keys()
                    if len(keys) == 1:
                        key = tuple(keys)[0]
                        msg = cls.schema.fields[key].error_messages['unique']
                        raise ValidationError({key: msg})
                    else:
                        fields = cls.schema.fields
                        # Compound index (sort value to make testing easier)
                        keys = sorted(keys)
                        raise ValidationError({k: fields[k].error_messages['unique_compound'].format(fields=keys)
                                               for k in keys})
            # Unknown index, cannot wrap the error so just reraise it
            raise
        return ret

    def io_validate(self, validate_all=False):
        """
        Run the io_validators of the document's fields.

        :param validate_all: If False only run the io_validators of the
            fields that have been modified.
        """
        if validate_all:
            return _io_validate_data_proxy(self.schema, self._data)
        else:
            return _io_validate_data_proxy(self.schema, self._data,
                                           partial=self._data.get_modified_fields_by_mongo_name())

    @classmethod
    async def find_one(cls, filter=None, *args, collection_cursor=False, **kwargs):
        """
        Find a single document in database.
        """
        filter = cook_find_filter(cls, filter)
        ret = await cls.collection.find_one(filter, *args, **kwargs)
        if not collection_cursor and ret is not None:
            ret = cls.build_from_mongo(ret, use_cls=True)
        return ret

    @classmethod
    def find(cls, filter=None, *args, collection_cursor=False, **kwargs):
        """
        Find a list document in database.

        Returns a cursor that provide Documents.
        """
        filter = cook_find_filter(cls, filter)
        return WrappedCursor(cls, cls.collection.find(filter, *args, **kwargs), collection_cursor)

    @classmethod
    async def find_one_and_update(cls, filter, update, *args, collection_cursor=False, **kwargs):
        """
        Finds a single document and updates it, returning either the
        original or the updated document.
        """
        filter = cook_find_filter(cls, filter)
        ret = await cls.collection.find_one_and_update(filter, update, *args, **kwargs)
        if not collection_cursor and ret is not None:
            ret = cls.build_from_mongo(ret, use_cls=True)
        return ret

    @classmethod
    async def count_documents(cls, filter=None, **kwargs):
        """
        Count the number of documents in this collection.
        """
        if filter is None:
            filter = {}
        filter = cook_find_filter(cls, filter)
        ret = await cls.collection.count_documents(filter, **kwargs)
        return ret

    @classmethod
    async def ensure_indexes(cls):
        """
        Check&create if needed the Document's indexes in database
        """
        for index in cls.opts.indexes:
            kwargs = index.document.copy()
            keys = [(k, d) for k, d in kwargs.pop('key').items()]
            await cls.collection.create_index(keys, **kwargs)


class MotorAsyncIOEmbeddedDocument(EmbeddedDocumentImplementation):

    __slots__ = ()

    opts = EmbeddedDocumentImplementation.opts

    async def reload(self):
        await self.parent_instance.reload()

    async def commit(self, *args, **kwargs):
        ret = await self.parent_instance.commit(*args, **kwargs)
        return ret

    async def delete(self, conditions=None):
        ret = await self.parent_instance.delete(conditions=conditions)
        return ret

    async def remove(self, conditions=None):
        ret = await self.parent_instance.remove(conditions=conditions)
        return ret

    def io_validate(self, validate_all=False):
        return self.parent_instance.io_validate(validate_all=validate_all)


# Run multiple validators and collect all errors in one
async def _run_validators(validators, field, value):
    errors = []
    tasks = [validator(field, value) for validator in validators]
    results = await asyncio.gather(*tasks, return_exceptions=True)
    for i, res in enumerate(results):
        if isinstance(res, ValidationError):
            errors.extend(res.messages)
        elif res:
            raise res
    if errors:
        raise ValidationError(errors)


async def _io_validate_data_proxy(schema, data_proxy, partial=None):
    errors = {}
    tasks = []
    tasks_field_name = []
    for name, field in schema.fields.items():
        data_name = field.attribute or name
        if partial is not None and data_name not in partial:
            continue

        value = data_proxy._data[data_name]
        if value is missing:
            continue
        try:
            if field.io_validate_recursive and not (value is None and field.allow_none):
                await field.io_validate_recursive(field, value)
            if field.io_validate:
                tasks.append(_run_validators(field.io_validate, field, value))
                tasks_field_name.append(name)
        except ValidationError as ve:
            errors[name] = ve.messages
    results = await asyncio.gather(*tasks, return_exceptions=True)
    for i, res in enumerate(results):
        if isinstance(res, ValidationError):
            errors[tasks_field_name[i]] = res.messages
        elif res:
            raise res
    if errors:
        raise ValidationError(errors)


async def _reference_io_validate(field, value):
    await value.fetch(no_data=True)


async def _list_io_validate(field, value):
    validators = field.container.io_validate
    if not validators or not value:
        return
    tasks = [_run_validators(validators, field.container, e) for e in value]
    results = await asyncio.gather(*tasks, return_exceptions=True)
    errors = {}
    for i, res in enumerate(results):
        if isinstance(res, ValidationError):
            errors[i] = res.messages
        elif res:
            raise res
    if errors:
        raise ValidationError(errors)


async def _embedded_document_io_validate(field, value):
    partial = value._data.get_modified_fields_by_mongo_name() if value.is_created else None
    await _io_validate_data_proxy(value.schema, value._data, partial=partial)


class MotorAsyncIOReference(Reference):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._document = None

    async def fetch(self, projection=None, no_data=False, force_reload=False):
        if not self._document or force_reload:
            if self.pk is None:
                raise ReferenceError('Cannot retrieve a None Reference')
            if no_data:
                result = await self.document_cls.count_documents({'_id': self.pk})
            else:
                result = self._document = await self.document_cls.find_one(self.pk, projection=projection)
            if not result:
                raise ValidationError(self.error_messages['not_found'].format(
                    document=self.document_cls.__name__))
        return self._document


class MotorAsyncIOBuilder(BaseBuilder):

    BASE_DOCUMENT_CLS = MotorAsyncIODocument
    BASE_EMBEDDED_DOCUMENT_CLS = MotorAsyncIOEmbeddedDocument

    @staticmethod
    def is_compatible_with(db):
        return isinstance(db, AsyncIOMotorDatabase)

    def _patch_field(self, field):
        super()._patch_field(field)

        validators = field.io_validate
        if not validators:
            field.io_validate = []
        else:
            if hasattr(validators, '__iter__'):
                validators = list(validators)
            else:
                validators = [validators]
            field.io_validate = [v if asyncio.iscoroutinefunction(v) else asyncio.coroutine(v)
                                 for v in validators]
        if isinstance(field, ListField):
            field.io_validate_recursive = _list_io_validate
        if isinstance(field, ReferenceField):
            field.io_validate.append(_reference_io_validate)
            field.reference_cls = MotorAsyncIOReference
        if isinstance(field, EmbeddedField):
            field.io_validate_recursive = _embedded_document_io_validate
