from mongoengine.errors import NotRegistered

__all__ = ('UPDATE_OPERATORS', 'get_document', '_document_registry')


UPDATE_OPERATORS = {'set', 'unset', 'inc', 'dec', 'mul',
                    'pop', 'push', 'push_all', 'pull',
                    'pull_all', 'add_to_set', 'set_on_insert',
                    'min', 'max', 'rename'}


_document_registry = {}


def get_document(name):
    """Get a registered Document class by name."""
    doc = _document_registry.get(name, None)
    if not doc:
        # Possible old style name
        single_end = name.split('.')[-1]
        compound_end = '.%s' % single_end
        possible_match = [k for k in _document_registry
                          if k.endswith(compound_end) or k == single_end]
        if len(possible_match) == 1:
            doc = _document_registry.get(possible_match.pop(), None)
    if not doc:
        raise NotRegistered("""
            `%s` has not been registered in the document registry.
            Importing the document class automatically registers it, has it
            been imported?
        """.strip() % name)
    return doc


def _get_documents_by_db(connection_alias, default_connection_alias):
    """Get all registered Documents class attached to a given database"""
    def get_doc_alias(doc_cls):
        return doc_cls._meta.get('db_alias', default_connection_alias)

    return [doc_cls for doc_cls in _document_registry.values()
            if get_doc_alias(doc_cls) == connection_alias]
