#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2023/4/15 11:05
# @Author  : old tom
# @File    : graph_helper.py
# @Project : futool-tiny-datahub
# @Desc    : 血缘关系上图，关系查询
from py2neo import Node, Relationship, Graph, Subgraph
from py2neo.matching import NodeMatcher


class Neo4jHelper(object):
    def __init__(self, graph: Graph):
        self.graph = graph
        self.node_matcher = NodeMatcher(graph)

    def find_node(self, *labels, **properties):
        """
        查找节点
        :param labels:
        :param properties:
        :return: 可能有多个
        """
        rt = self.node_matcher.match(*labels, **properties)
        return rt

    def find_node_by_id(self, node_id):
        """
        节点ID 查找节点
        :param node_id:
        :return:
        """
        return self.node_matcher.get(node_id)

    def check_node_exists(self, *labels, **properties):
        """
        判断节点是否存在
        :param labels:
        :param properties:
        :return:
        """
        return self.find_node(*labels, **properties).exists()

    def same_node_count(self, *labels, **properties):
        """
        相同节点数
        :param labels:
        :param properties:
        :return:
        """
        return self.find_node(labels, properties).count()

    def find_same_node(self, *labels, **properties):
        """
        查找相同节点
        :param labels:
        :param properties:
        :return:
        """
        return self.find_node(*labels, **properties).all()

    def merge_subgraph(self, relationship, label, *property_keys):
        """
        合并子图
        :param property_keys: 判断节点重复属性
        :param label: 节点标签
        :param relationship:
        :return:
        """
        self.graph.merge(relationship, label, *property_keys)

    def find_subgraph(self):
        pass

    def create_subgraph(self, nodes: list, edges: list):
        """
        创建子图,注：相同节点不会自动去重
        :param nodes: 节点列表
        :param edges: 边（关系列表）
        :return:
        """
        # 开启事务
        tx = self.graph.begin()
        try:
            tx.create(Subgraph(nodes=nodes, relationships=edges))
            self.graph.commit(tx)
        except Exception as e:
            self.graph.rollback(tx)

    def create_node(self, *labels, **properties):
        """
        创建节点
        :param node:
        :return:
        """
        node = Node(*labels, **properties)
        self.graph.create(node)
        return node

    def create_node_with_check(self, *labels, **properties):
        """
        创建节点前判断是否存在
        :param labels:
        :param properties:
        :return: 创建的节点或已经存在的节点
        """
        if not self.check_node_exists(*labels, **properties):
            return self.create_node(*labels, **properties)
        return self.find_node(*labels, **properties).first()

    def create_relationship(self, from_node, relation_desc, to_node):
        """
        创建关系
        :param to_node:
        :param relation_desc: 关系描述
        :param from_node:
        :return:
        """
        self.graph.create(Relationship(from_node, relation_desc, to_node))

    def create_relationship_with_merge(self, from_node, relation_desc, to_node, label, *property_keys):
        """
        创建关系并合并相同节点
        :param label: 标签
        :param from_node:
        :param relation_desc: 关系描述
        :param to_node:
        :return:
        """
        relation = Relationship(from_node, relation_desc, to_node)
        self.graph.create(relation)
        self.merge_subgraph(relation, label, *property_keys)