import copy
import datetime
import itertools
import json
import re
import time
import functools

from elasticsearch import (Elasticsearch, NotFoundError, RequestError,
                           TransportError, ElasticsearchException)
from elasticsearch import helpers
from elasticsearch.serializer import JSONSerializer
from importlib import import_module

from biothings.utils.common import iter_n, splitstr, nan, inf, merge, traverse
from biothings.utils.dataload import dict_walk
from biothings.utils.serializer import to_json

# setup ES logging
import logging
formatter = logging.Formatter("%(levelname)s:%(name)s:%(message)s")
es_logger = logging.getLogger('elasticsearch')
es_logger.setLevel(logging.WARNING)
ch = logging.StreamHandler()
ch.setFormatter(formatter)
es_logger.addHandler(ch)

es_tracer = logging.getLogger('elasticsearch.trace')
es_tracer.setLevel(logging.WARNING)
ch = logging.StreamHandler()
ch.setFormatter(formatter)
es_tracer.addHandler(ch)

def verify_ids(doc_iter, es_host, index, doc_type=None, step=100000, ):
    '''verify how many docs from input interator/list overlapping with existing docs.'''

    index = index
    doc_type = doc_type
    es = get_es(es_host)
    q = {'query': {'ids': {"values": []}}}
    total_cnt = 0
    found_cnt = 0
    out = []
    for doc_batch in iter_n(doc_iter, n=step):
        id_li = [doc['_id'] for doc in doc_batch]
        # id_li = [doc['_id'].replace('chr', '') for doc in doc_batch]
        q['query']['ids']['values'] = id_li
        xres = es.search(index=index, doc_type=doc_type, body=q, _source=False)
        found_cnt += xres['hits']['total']
        total_cnt += len(id_li)
        out.extend([x['_id'] for x in xres['hits']['hits']])
    return out


def get_es(es_host, timeout=120, max_retries=3, retry_on_timeout=False):
    es = Elasticsearch(es_host, timeout=timeout, max_retries=max_retries,
                       retry_on_timeout=retry_on_timeout)
    return es


def wrapper(func):
    '''this wrapper allows passing index and doc_type from wrapped method.'''
    def outter_fn(*args, **kwargs):
        self = args[0]
        index = kwargs.pop('index', self._index)             # pylint: disable=protected-access
        doc_type = kwargs.pop('doc_type', self._doc_type)    # pylint: disable=protected-access
        self._index = index                                  # pylint: disable=protected-access
        self._doc_type = doc_type                            # pylint: disable=protected-access
        return func(*args, **kwargs)
    outter_fn.__doc__ = func.__doc__
    return outter_fn


class IndexerException(Exception):
    pass

class ESIndex():
    """ An Elasticsearch Index Wrapping A Client.
    Counterpart for pymongo.collection.Collection """

    # a new approach to biothings.utils.es.ESIndexer
    # but not intended to be a replacement in features.

    def __init__(self, client, index_name):
        self.client = client  # synchronous
        self.index_name
 = index_name  # MUST exist

    @property
    @functools.lru_cache()
    def doc_type(self):
        if int(self.client.info()['version']['number'].split('.')[0]) < 7:
            mappings = self.client.indices.get_mapping(self.index_name)
            mappings = mappings[self.index_name]["mappings"]
            return next(iter(mappings.keys()))
        return None

    # SUBCLASS NOTE &&
    # BEFORE YOU AD
D A METHOD UNDER THIS CLASS:

    # An object of this class refers to an existing ES index. All operations
    # should target this index. Do not put uncommon methods here. They belong
    # to the subclasses. This class is to provide a common framework to support
    # Index-specific ES operations, an concept does not exist in the low-level
    # ES library, thus only providing low-level common operations, like doc_type
    # parsing across ES versions for biothings usecases. (single type per index)

class ESIndexer():
    # RETIRING THIS CLASS

    # --
    # Since we don't always directly work on
    # existing actual indices, the index referred here
    # can be an alias or does not exist.

    def __init__(self, index, doc_type='_doc', es_host='localhost:9200',
                 step=500, step_size=10,  # elasticsearch.helpers.bulk
                 number_of_shards=1, number_of_replicas=0,
                 check_index=True, **kwargs):
        self.es_host = es_host
        self._es = get_es(es_host, **kwargs)
        self._host_major_ver = int(self._es.info()['version']['number'].split('.')[0])
        if check_index:
            # if index is actually an alias, resolve the alias to
            # the real underlying index
            try:
                res = self._es.indices.get_alias(index=index)
                # this was an alias
                assert len(res) == 1, "Expecing '%s' to be an alias, but got nothing..." % index
                self._index = list(res.keys())[0]
            except NotFoundError:
                # this was a real index name
                self._index = index
        self._doc_type = None
        if doc_type:
            self._doc_type = doc_type
        else:
            # assuming index exists, get mapping to discover doc_type
            try:
                m = self.get_mapping()
                assert len(m) == 1, "Expected only one doc type, got: %s" % m.keys()
                self._doc_type = list(m).pop()
            except Exception as e:       # pylint: disable=broad-except
                if check_index:
                    logging.info("Failed to guess doc_type: %s", e)
        self.number_of_shards = number_of_shards            # set number_of_shards when create_index
        self.number_of_replicas = int(number_of_replicas)   # set number_of_replicas when create_index
        self.step = step or 500   # the bulk size when doing bulk operation.
        self.step_size = (step_size or 10) * 1048576  # MB -> bytes
        self.s = None      # number of records to skip, useful to continue indexing after an error.

    @wrapper
    def get_biothing(self, bid, only_source=False, **kwargs):
        rawdoc = self._es.get(index=self._index, id=bid, doc_type=self._doc_type, **kwargs)
        if not only_source:
            return rawdoc
        else:
            doc = rawdoc['_source']
            doc["_id"] = rawdoc["_id"]
            return doc

    @wrapper
    def exists(self, bid):
        """return True/False if a biothing id exists or not."""
        try:
            doc = self.get_biothing(bid, stored_fields=None)
            return doc['found']
        except NotFoundError:
            return False

    @wrapper
    def mexists(self, bid_list):
        q = {
            "query": {
                "ids": {
                    "values": bid_list
                }
            }
        }
        res = self._es.search(index=self._index, doc_type=self._doc_type, body=q, stored_fields=None, size=len(bid_list))
        # id_set = set([doc['_id'] for doc in res['hits']['hits']])     # TODO: Confirm this line
        id_set = {doc['_id'] for doc in res['hits']['hits']}
        return [(bid, bid in id_set) for bid in bid_list]

    @wrapper
    def count(self, q=None, raw=False):
        try:
            _res = self._es.count(index=self._index, doc_type=self._doc_type, body=q)
            return _res if raw else _res['count']
        except NotFoundError:
            return None

    @wrapper
    def count_src(self, src):
        if isinstance(src, str):
            src = [src]
        cnt_d = {}
        for _src in src:
            q = {
                "query": {
                    "constant_score": {
                        "filter": {
                            "exists": {"field": _src}
                        }
                    }
                }
            }
            cnt_d[_src] = self.count(q)
        return cnt_d

    @wrapper
    def create_index(self, mapping=None, extra_settings=None):
        if not self._es.indices.exists(index=self._index):
            body = {
                'settings': {
                    'number_of_shards': self.number_of_shards,
                    "number_of_replicas": self.number_of_replicas,
                }