# coding: utf-8

"""
Support for Elasticsearch
-------------------------

Provides a couple of template tasks to use with Elasticsearch indexing.
"""

from elasticsearch.helpers import bulk_index
import abc
import elasticsearch
import json
import luigi

MARC_MAPPING = {'title': {'date_detection': False,
                          '_id': {'path': 'content.001'},
                          '_all': {'enabled': True,
                                   'term_vector': 'with_positions_offsets',
                                   'store': True}}}


class ESIndexTask(luigi.Task):
    """
    Template task for indexing data with Elasticsearch.

    Usage:
    Subclass and override the required ``index`` and ``doc_type`` attributes.

    To customize how to access data from an input task, override the `docs` method
    with a generator that yields each document as a json string.

    Example:

        import luigi
        from tasktree.elastic import ElasticsearchIndex

        class MyIndex(ElasticsearchIndex):
            index = 'example'
            doc_type = 'default'

            def docs(self):
                for i in range(10):
                    yield { '_index': 'example',
                            '_type': 'default',
                            '_id': 'doc-%s' % i.
                            '_source': {'text': 'Hello World-%s' % i}}

        if __name__ == '__main__':
            luigi.run()

    """
    hosts = [{'host': 'localhost', 'port': 9200}]
    index = NotImplemented
    doc_type = NotImplemented
    settings = {}
    mapping = {}

    def exists(self):
        """
        Return ``True`` if an index with the given name exists.
        """
        es = elasticsearch.Elasticsearch(hosts=self.hosts)
        return es.indices.exists(self.index)

    def expected(self):
        """
        The **expected** number of docs in the index. Override,
        if you have duplicates in your docs (e.g. same id) - but better,
        remove duplicates before indexing.

        .. admonition:: Slow default implementation

            Also, the default implementation might take too long,
            since it has to iterate over the whole `self.docs` once.
        """
        return 0 # len(list(self.docs()))

    def indexed(self):
        """
        Return the number of indexed documents for this index
        and doc_type.
        """
        es = elasticsearch.Elasticsearch(hosts=self.hosts)
        result = es.count(index=self.index, doc_type=self.doc_type)
        count = result.get('count', -1)
        return count

    @abc.abstractmethod
    def docs(self):
        """
        Return an iterable with the documents to be indexed.
        Implement this in you subclass or use some prepared
        subclasses for the common cases.
        """
        return []

    def complete(self):
        """
        Override the complete condition here, if you have special needs.

        By default an index is considered complete, if:

        * it exists and
        * the number of documents in the index equals the number
          of documents, the iterator yields

        Stronger conditions might be in order, as can be seen in
        :class:`ValueComparatorIndex`.
        """
        if not self.exists():
            return False
        if not self.expected() == self.indexed():
            return False
        return True

    def representative(self):
        """
        Return an example document (particularly its `_source`),
        but do not care which.

        Raises a ``RuntimeError`` if no document can be retrieved at all.
        """
        es = elasticsearch.Elasticsearch(hosts=self.hosts)
        result = es.search(index=self.index, body={'query': {'match_all' : {}}},
                           size=1)
        hits = result.get('hits').get('hits')
        if len(hits) == 0:
            raise RuntimeError('No docs in index: %s' % self.index)
        return hits[0].get('_source')

    def run(self):
        """
        (Bulk) Index data generated by docs into the target index.

        The default, hard-coded chunk size is 2000 docs per request.

        If `load_balance` is trueish, the number of replicas will be set to
        zero during indexing time, essentially distributing the indexing work
        among the nodes. After indexing the replicas set back to the default 1.

        When the index already exists, it will be deleted first.
        """
        es = elasticsearch.Elasticsearch(hosts=self.hosts)
        if es.indices.exists(self.index):
            es.indices.delete(index=self.index)
        es.indices.create(index=self.index)
        if self.mapping:
            es.indices.put_mapping(index=self.index, doc_type=self.doc_type,
                                   body=self.mapping)

        es.indices.put_settings({"index": {"refresh_interval": "-1"}},
                                index=self.index)

        bulk_index(es, self.docs(), chunk_size=2000, raise_on_error=True)

        es.indices.put_settings({"index": {"refresh_interval": "1s"}},
                                index=self.index)


class DocIndexTask(ESIndexTask):
    """
    Implements a default iterator over the docs.

    Assumptions:

    1. There is only **one** input to the task.
    2. The input is a JSON file, with one document per line (.ndj).
    3. The id of the document can be found in ``content.001``.
    """

    def docs(self):
        """
        A minimal implementation of docs. This method needs to return an
        iterable of dict objects, each with the following keys: `_index`,
        `_type`, `_id` and `_source`. This is a bit redundant, since index and
        doc_type are already known, but it is the API of the `bulk_index` helper.
        """
        with self.input().open() as handle:
            documents = ({
                '_index': self.index,
                '_type': self.doc_type,
                '_id': doc.get('content').get('001'),
                '_source': doc} for doc in (json.loads(line) for line in handle))
            for doc in documents:
                yield doc
