# coding=utf-8
import elasticsearch
from elasticsearch import Elasticsearch as Es
from django.conf import settings
from pytz import timezone
from libs.debug import pretty_json
import logging
import traceback
import py2neo
from py2neo import Graph
import json
from libs.error import logging_exception


class Neo4jManager(object):
    graph_obj = None

    @classmethod
    def neo4j_get_graph_obj(cls):
        try:
            if not cls.graph_obj:
                cls.graph_obj = Graph(host=settings.NEO4J_HOST, auth=(settings.NEO4J_USER, settings.NEO4J_PWD))

            return cls.graph_obj
        except:
            logging.error("catch exception,logins:%s" % traceback.format_exc())
            return None

    @classmethod
    def neo4j_get_related_word_by_graph(cls, graph_obj=None, query_word="", graph_related_word_set=set(),
                                        have_read_name_set=set(), need_size=6, redis_client=None,need_many_relation=False):
        try:
            assert (query_word and len(query_word) > 0)
            assert (graph_obj is not None)
            assert (redis_client is not None)

            query_forbidden_redis_key = "doris_search:neo4j_forbidden_word_list:word:" + str(query_word)
            query_forbidden_redis_data = redis_client.get(query_forbidden_redis_key)

            query_forbidden_word_list = list()
            if query_forbidden_redis_data:
                query_forbidden_word_list = json.loads(str(query_forbidden_redis_data, encoding="utf-8"))

            query_neo4j_redis_key = "doris_search:query_neo4j_word:" + str(query_word)
            query_neo4j_redis_data = redis_client.get(query_neo4j_redis_key)

            query_neo4j_word_set = set()
            if query_neo4j_redis_data:
                query_neo4j_word_set = set(json.loads(str(query_neo4j_redis_data,encoding="utf-8")))
            else:
                # if need_many_relation:
                #     graph_query = "match (na)-[re]->(nb) where na.name = \"" + query_word + "\" WITH na,re,nb match (nb)-[re2]->(nc) return na,re,nb,re2,nc"
                # else:
                graph_query = "match (na)-[re]->(nb) where na.name = \"" + query_word + "\" return na,re,nb"
                graph_results = graph_obj.run(graph_query)

                graph_total_num = 0
                graph_valid_num = 0
                for item in graph_results:
                    nb = item.get("nb").get("name")
                    graph_total_num += 1
                    if nb != query_word:
                        query_neo4j_word_set.add(nb)
                        graph_valid_num += 1

                    if need_many_relation:
                        nc = item.get("nc").get("name")
                        graph_total_num += 1
                        if nc != query_word:
                            query_neo4j_word_set.add(nc)
                            graph_valid_num += 1

                redis_client.set(query_neo4j_redis_key,json.dumps(list(query_neo4j_word_set)))
                redis_client.expire(query_neo4j_redis_key,2*24*60*60)

            valid_query_neo4j_word_set = query_neo4j_word_set - have_read_name_set - set(query_forbidden_word_list) - graph_related_word_set

            get_neo4j_word_num = 0
            for word in valid_query_neo4j_word_set:
                graph_related_word_set.add(word)
                get_neo4j_word_num += 1
                if get_neo4j_word_num >= need_size:
                    break

            logging.info("query_word:%s graph related word num:%d" % (query_word,len(query_neo4j_word_set)))

            return graph_related_word_set
        except:
            logging_exception()
            logging.error("catch exception,err_msg:%s" % traceback.format_exc())
            return graph_related_word_set
