# -*- coding: utf-8 -*-
"""
NOTE: You should use the unit tests, not these doctests, which are harder to get running consistently.
I've left them here as documentation only, they are accurate as usage examples.

Create ElasticSearch connection
>>> conn = ElasticSearch('http://localhost:9200/')

Or a more verbose log level.
>>> import logging
>>> class VerboseElasticSearch(ElasticSearch):
...     def setup_logging(self):
...         log = super(VerboseElasticSearch, self).setup_logging()
...         log.addHandler(logging.StreamHandler())
...         log.setLevel(logging.DEBUG)
...         return log
>>> conn = VerboseElasticSearch('http://localhost:9200/')

Add a few documents

>>> conn.index({"name":"Joe Tester"}, "test-index", "test-type", 1)
{'_type': 'test-type', '_id': '1', 'ok': True, '_index': 'test-index'}
>>> conn.index({"name":"Bill Baloney"}, "test-index", "test-type", 2)
{'_type': 'test-type', '_id': '2', 'ok': True, '_index': 'test-index'}

Get one

>>> conn.refresh(["test-index"]) # doctest: +ELLIPSIS
{'ok': True, '_shards': {...}}
>>> conn.get("test-index", "test-type", 1)
{'_type': 'test-type', '_id': '1', '_source': {'name': 'Joe Tester'}, '_index': 'test-index'}

Get a count
>>> conn.count("name:joe")
{'count': 1, '_shards': {'successful': 5, 'failed': 0, 'total': 5}}

Search

>>> conn.search("name:joe")
{'hits': {'hits': [{'_type': 'test-type', '_id': '1', '_source': {'name': 'Joe Tester'}, '_index': 'test-index'}], 'total': 1}, '_shards': {'successful': 5, 'failed': 0, 'total': 5}}

Terms

>>> conn.terms(['name'])
{'docs': {'max_doc': 2, 'num_docs': 2, 'deleted_docs': 0}, 'fields': {'name': {'terms': [{'term': 'baloney', 'doc_freq': 1}, {'term': 'bill', 'doc_freq': 1}, {'term': 'joe', 'doc_freq': 1}, {'term': 'tester', 'doc_freq': 1}]}}, '_shards': {'successful': 5, 'failed': 0, 'total': 5}}
>>> conn.terms(['name'], indexes=['test-index'])
{'docs': {'max_doc': 2, 'num_docs': 2, 'deleted_docs': 0}, 'fields': {'name': {'terms': [{'term': 'baloney', 'doc_freq': 1}, {'term': 'bill', 'doc_freq': 1}, {'term': 'joe', 'doc_freq': 1}, {'term': 'tester', 'doc_freq': 1}]}}, '_shards': {'successful': 5, 'failed': 0, 'total': 5}}
>>> conn.terms(['name'], min_freq=2)
{'docs': {'max_doc': 2, 'num_docs': 2, 'deleted_docs': 0}, 'fields': {'name': {'terms': []}}, '_shards': {'successful': 5, 'failed': 0, 'total': 5}}

More Like This
>>> conn.index({"name":"Joe Test"}, "test-index", "test-type", 3)
{'_type': 'test-type', '_id': '3', 'ok': True, '_index': 'test-index'}
>>> conn.refresh(["test-index"]) # doctest: +ELLIPSIS
{'ok': True, '_shards': {...}}
>>> conn.morelikethis("test-index", "test-type", 1, ['name'], min_term_freq=1, min_doc_freq=1)
{'hits': {'hits': [{'_type': 'test-type', '_id': '3', '_source': {'name': 'Joe Test'}, '_index': 'test-index'}], 'total': 1}, '_shards': {'successful': 5, 'failed': 0, 'total': 5}}
>>> conn.delete("test-index", "test-type", 3)
{'_type': 'test-type', '_id': '3', 'ok': True, '_index': 'test-index'}

Delete Bill

>>> conn.delete("test-index", "test-type", 2)
{'_type': 'test-type', '_id': '2', 'ok': True, '_index': 'test-index'}

>>> conn.delete_by_query("test-index, "test-type", {"query_string": {"query": "name:joe OR name:bill"}})
{'ok': True, '_indices': {'test-index': {'_shards': {'successful': 5, 'failed': 0, 'total': 5}}}}

Delete the index

>>> conn.delete_index("test-index")
{'acknowledged': True, 'ok': True}

Create the index anew

>>> conn.create_index("test-index")
{'acknowledged': True, 'ok': True}

Try (and fail) to create an existing index

>>> conn.create_index("test-index")
{'error': '[test-index] Already exists'}

Put mapping

>>> conn.put_mapping("test-type", {"test-type" : {"properties" : {"name" : {"type" : "string", "store" : "yes"}}}})
{'acknowledged': True, 'ok': True}

Get status

>>> conn.status(["test-index"]) # doctest: +ELLIPSIS
{'indices': {'test-index': ...}}

>>> conn.flush(["test-index"]) # doctest: +ELLIPSIS
{'ok': True, '_shards': {...}}

>>> conn.refresh(["test-index"]) # doctest: +ELLIPSIS
{'ok': True, '_shards': {...}}

>>> conn.optimize(["test-index"]) # doctest: +ELLIPSIS
{'ok': True, '_shards': {...}}

Test adding with automatic id generation
>>> conn.index({"name":"Joe Tester"}, "test-index", "test-type") # doctest: +ELLIPSIS
{'_type': 'test-type', '_id': '...', 'ok': True, '_index': 'test-index'}


"""

__author__ = 'Robert Eanes'
__all__ = ['ElasticSearch']
__version__ = (0, 0, 3)

def get_version():
    return "%s.%s.%s" % __version__

try:
    # For Python < 2.6 or people using a newer version of simplejson
    import simplejson as json
except ImportError:
    # For Python >= 2.6
    import json

from urllib import urlencode
import logging
import requests


class ElasticSearchError(Exception):
    pass



class NullHandler(logging.Handler):
    def emit(self, record):
        pass



class ElasticSearch(object):
    """
    ElasticSearch connection object.
    """

    def __init__(self, url, timeout=60):
        self.url = url
        self.timeout = timeout

        if self.url.endswith('/'):
            self.url = self.url[:-1]

    def setup_logging(self):
        """
        Sets up the logging.

        Done as a method so others can override as needed without complex
        setup.
        """
        log = logging.getLogger('pyelasticsearch')
        null = NullHandler()
        log.addHandler(null)
        log.setLevel(logging.ERROR)
        return log

    def _make_path(self, path_components):
        """
        Smush together the path components. Empty components will be ignored.
        """
        path_components = [str(component) for component in path_components if component]
        path = '/'.join(path_components)
        if not path.startswith('/'):
            path = '/'+path
        return path

    def _build_url(self, path):
        return self.url + path

    def _send_request(self, method, path, body="", querystring_args={}):
        if querystring_args:
            path = "?".join([path, urlencode(querystring_args)])

        kwargs = {
            'timeout': self.timeout,
        }
        url = self._build_url(path)

        if body:
            kwargs['data'] = self._prep_request(body)

        if not hasattr(requests, method.lower()):
            raise ElasticSearchError("No such HTTP Method '%s'!" % method.lower())

        logging.debug("making %s request to path: %s %s with body: %s" % (method, url, path, kwargs.get('data', {})))
        req_method = getattr(requests, method.lower())
        resp = req_method(url, **kwargs)
        logging.debug("response status: %s" % resp.status_code)
        prepped_response = self._prep_response(resp.content)
        logging.debug("got response %s" % prepped_response)
        return prepped_response

    def _prep_request(self, body):
        """
        Encodes body as json.
        """
        return json.dumps(body)

    def _prep_response(self, response):
        """
        Parses json to a native python object.
        """
        return json.loads(response)

    def _query_call(self, query_type, query, body=None, indexes=['_all'], doc_types=[], **query_params):
        """
        This can be used for search and count calls.
        These are identical api calls, except for the type of query.
        """
        querystring_args = query_params
        if query:
            querystring_args['q'] = query
        path = self._make_path([','.join(indexes), ','.join(doc_types),query_type])
        response = self._send_request('GET', path, body, querystring_args)
        return response

    ## REST API

    def index(self, doc, index, doc_type, id=None, force_insert=False):
        """
    	Index a typed JSON document into a specific index and make it searchable.
        """
        if force_insert:
            querystring_args = {'op_type':'create'}
        else:
            querystring_args = {}

        if id is None:
            request_method = 'POST'
        else:
            request_method = 'PUT'
        path = self._make_path([index, doc_type, id])
        response = self._send_request(request_method, path, doc, querystring_args)
        return response

    def delete(self, index, doc_type, id):
        """
        Delete a typed JSON document from a specific index based on its id.
        """
        path = self._make_path([index, doc_type, id])
        response = self._send_request('DELETE', path)
        return response

    def delete_by_query(self, index, doc_type, query):
        """
        Delete a typed JSON documents from a specific index based on query
        """
        path = self._make_path([index, doc_type, '_query'])
        response = self._send_request('DELETE', path, query)
        return response

    def get(self, index, doc_type, id):
        """
        Get a typed JSON document from an index based on its id.
        """
        path = self._make_path([index, doc_type, id])
        response = self._send_request('GET', path)
        return response

    def search(self, query, body=None, indexes=['_all'], doc_types=[], **query_params):
        """
        Execute a search query against one or more indices and get back search hits.
        query must be a dictionary that will convert to Query DSL
        TODO: better api to reflect that the query can be either 'query' or 'body' argument.
        """
        return self._query_call("_search", query, body, indexes, doc_types, **query_params)

    def count(self, query, body=None, indexes=['_all'], doc_types=[], **query_params):
        """
        Execute a query against one or more indices and get hits count.
        """
        return self._query_call("_count", query, body, indexes, doc_types, **query_params)

    def put_mapping(self, doc_type, mapping, indexes=['_all']):
        """
        Register specific mapping definition for a specific type against one or more indices.
        """
        path = self._make_path([','.join(indexes), doc_type,"_mapping"])
        response = self._send_request('PUT', path, mapping)
        return response

    def morelikethis(self, index, doc_type, id, fields, **query_params):
        """
        Execute a "more like this" search query against one or more fields and get back search hits.
        """
        path = self._make_path([index, doc_type, id, '_mlt'])
        query_params['fields'] = ','.join(fields)
        response = self._send_request('GET', path, querystring_args=query_params)
        return response

    ## Index Admin API

    def status(self, indexes=['_all']):
        """
        Retrieve the status of one or more indices
        """
        path = self._make_path([','.join(indexes), '_status'])
        response = self._send_request('GET', path)
        return response

    def create_index(self, index, settings=None):
        """
        Creates an index with optional settings.
        Settings must be a dictionary which will be converted to JSON.
        Elasticsearch also accepts yaml, but we are only passing JSON.
        """
        response = self._send_request('PUT', self._make_path([index]), settings)
        return response

    def delete_index(self, index):
        """
        Deletes an index.
        """
        response = self._send_request('DELETE', self._make_path([index]))
        return response

    def flush(self, indexes=['_all'], refresh=None):
        """
        Flushes one or more indices (clear memory)
        """
        path = self._make_path([','.join(indexes), '_flush'])
        args = {}
        if refresh is not None:
            args['refresh'] = refresh
        response = self._send_request('POST', path, querystring_args=args)
        return response

    def refresh(self, indexes=['_all']):
        """
        Refresh one or more indices
        """
        path = self._make_path([','.join(indexes), '_refresh'])
        response = self._send_request('POST', path)
        return response

    def gateway_snapshot(self, indexes=['_all']):
        """
        Gateway snapshot one or more indices
        """
        path = self._make_path([','.join(indexes), '_gateway', 'snapshot'])
        response = self._send_request('POST', path)
        return response


    def optimize(self, indexes=['_all'], **args):
        """
        Optimize one ore more indices
        """
        path = self._make_path([','.join(indexes), '_optimize'])
        response = self._send_request('POST', path, querystring_args=args)
        return response

if __name__ == "__main__":
    import doctest
    doctest.testmod()
