from elasticsearch import Elasticsearch as Es
from elasticsearch import helpers


def es_index_adapt(index_prefix, doc_type, rw=None):
    """get the adapted index name
    """
    assert rw in [None, "read", "write"]
    index = "-".join((index_prefix, doc_type))
    if rw:
        index = "-".join((index, rw))
    return index


def get_es():
    init_args = {
        "sniff_on_start": False,
        "sniff_on_connection_fail": False,
    }
    new_hosts = [{
        "host": "172.16.31.17",
        "port": 9000,
    }, {
        "host": "172.16.31.11",
        "port": 9000,
    }, {
        "host": "172.16.31.13",
        "port": 9000,
    }]
    new_es = Es(hosts=new_hosts, **init_args)
    return new_es


def es_query(doc, body, offset=0, size=100, es=None, rw="read"):
    if es is None:
        es = get_es()
    index = es_index_adapt(index_prefix="gm-dbmw", doc_type=doc, rw=rw)
    res = es.search(index=index, timeout="10s", body=body, from_=offset, size=size)
    return res


def es_scan(doc, body, es=None, rw="read"):
    if es is None:
        es = get_es()
    index = es_index_adapt(index_prefix="gm-dbmw", doc_type=doc, rw=rw)
    return helpers.scan(es, index=index, query=body, request_timeout=100, scroll="300m", raise_on_error=False)


# curl "http://172.16.31.17:9000/gm-dbmw-tractate-read/_search?pretty&size=1" -d '
# '
def get_tractate_info_from_es(fields=["id"]):
    q = {
        "query": {
            "bool": {
                "must": [{
                    "term": {
                        "is_online": "true"
                    }
                }, {
                    "terms": {
                        "content_level": [6, 5, 4, 3.5, 3]
                    }
                }],
                "must_not": [{
                    "term": {
                        "status": 4
                    }
                }, {
                    "term": {
                        "show_by_index": 2
                    }
                }]
            }
        },
        "_source": {
            "include": fields
        }
    }
    results = es_scan("tractate", q)
    return results


def get_answer_info_from_es(fields=["id"]):
    q = {
        "query": {
            "bool": {
                "must": [{
                    "term": {
                        "is_online": True
                    }
                }, {
                    "terms": {
                        "content_level": [6, 5, 4, 3.5, 3]
                    }
                }, {
                    "range": {
                        "content_length": {
                            "gte": 30
                        }
                    }
                }]
            }
        },
        "_source": {
            "include": fields
        }
    }
    results = es_scan("answer", q)
    return results


def get_diary_info_from_es(fields=["id"]):
    q = {
        "query": {
            "bool": {
                "filter": [{
                    "term": {
                        "is_online": True
                    }
                }, {
                    "term": {
                        "has_cover": True
                    }
                }, {
                    "term": {
                        "is_sink": False
                    }
                }, {
                    "term": {
                        "has_before_cover": True
                    }
                }, {
                    "term": {
                        "has_after_cover": True
                    }
                }, {
                    "terms": {
                        "content_level": [6, 5, 4, 3.5, 3]
                    }
                }, {
                    "term": {
                        "content_simi_bol_show": 0
                    }
                }, {
                    "exists": {
                        "field": "before_cover_url"
                    }
                }]
            }
        },
        "_source": {
            "include": fields
        }
    }
    results = es_scan("diary", q)
    return results
