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 save_diary_info_from_es():
#     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": ["id"]
#         }
#     }
#     count = 0
#     # before_res_dict = {}
#     after_res_dict = {}
#     results = es_scan("diary", q)
#     for item in results:
#         diary_id = item["_id"]
#         # before_cover_url = item["_source"]["before_cover_url"] + "-w"
#         # before_img = url_to_ndarray(before_cover_url)

#         after_cover_url = item["_source"]["after_cover_url"] + "-w"
#         img = url_to_ndarray(after_cover_url)
#         if img.any():
#             count += 1
#             print("count: " + str(count) + " " + str(diary_id))
#             faces = FACE_TO_VEC_FUN(img)
#             for face in faces:
#                 after_res_dict[diary_id] = face["feature"]
#     redis_client_db.hmset(DIARY_AFTER_COVER_FEATURE_KEY, after_res_dict)
