from neo4j import GraphDatabase
from common.config import Config
from tqdm import tqdm
import json

conf = Config()


class Neo4jClient:
    def __init__(self, uri, user, password):
        """初始化连接"""
        self.driver = GraphDatabase.driver(uri, auth=(user, password))

    def __del__(self):
        """关闭连接"""
        if self.driver is not None:
            self.driver.close()

    def run_cypher(self, query, parameters=None):
        """
        执行一条 Cypher 语句并返回结果
        :param query: Cypher 查询语句
        :param parameters: 可选参数字典
        :return: 查询结果列表（每一行是一个 dict）
        """
        with self.driver.session() as session:
            result = session.run(query, parameters or {})
            return [record.data() for record in result]

    def run_multiple_cypher(self, queries_with_params):
        """
        在一个事务里执行多条 Cypher 语句，显示 tqdm 进度条，返回每条语句的结果。

        整批是一个事务：任意一条报错，整批回滚，不会写进去一半。
        批量导入图谱时用这个，比一条条 run_cypher 快很多（少了每条一次的往返和提交）。

        参数:
            queries_with_params: List[Tuple[str, Dict]]
                形式如: [("CREATE (n:Test {name: $name})", {"name": "Alice"}), ...]
                参数字典可以传 None，等价于空字典。
        返回:
            List[List[Dict]] —— 与入参一一对应，每项是那条语句的结果行。
            CREATE/MERGE 这类没有 RETURN 的语句，对应项是空列表 []。
        """
        with self.driver.session() as session:

            def transaction_logic(tx):
                results = []
                for query, params in tqdm(queries_with_params, desc="执行 Cypher 语句"):
                    result = tx.run(query, params or {})
                    # 必须在事务里就把记录读出来：Result 是惰性的，事务一提交就失效，
                    # 拿到外面再遍历会抛 ResultConsumedError。
                    results.append([record.data() for record in result])
                return results

            return session.execute_write(transaction_logic)

    def export_tcm_metadata_to_json(self, output_path="tcm_metadata.json"):
        with self.driver.session() as session:

            # 1. 所有节点标签
            label_query = """
            MATCH (n)
            UNWIND labels(n) AS label
            RETURN DISTINCT label
            """
            labels = [record["label"] for record in session.run(label_query)]

            # 2. 所有关系类型
            rel_query = """
            MATCH (n)-[r]-()
            RETURN DISTINCT type(r) AS rel_type
            """
            rel_types = [record["rel_type"] for record in session.run(rel_query)]

            # 3. 所有三元组结构
            triple_query = """
            MATCH (n)-[r]->(m)
            WITH head(labels(n)) AS from_label, type(r) AS rel_type, head(labels(m)) AS to_label
            RETURN DISTINCT from_label, rel_type, to_label
            """
            triples = [{
                "from": record["from_label"],
                "rel_type": record["rel_type"],
                "to": record["to_label"],
                "description": ""
            } for record in session.run(triple_query)]

            # 4. 节点属性（每个标签下的属性键）
            node_props_query = """
            MATCH (n)
            UNWIND labels(n) AS label
            UNWIND keys(n) AS prop
            RETURN DISTINCT label, prop
            ORDER BY label, prop
            """
            label_props = {}
            for record in session.run(node_props_query):
                label = record["label"]
                prop = record["prop"]
                if prop == "project":  # 忽略 project 字段
                    continue
                label_props.setdefault(label, []).append({
                    "name": prop,
                    "description": ""
                })

            # 5. 关系属性（每种关系下的属性键）
            rel_props_query = """
            MATCH (n)-[r]->(m)
            UNWIND keys(r) AS prop
            RETURN DISTINCT type(r) AS rel_type, prop
            ORDER BY rel_type, prop
            """
            rel_type_props = {}
            for record in session.run(rel_props_query):
                rel_type = record["rel_type"]
                prop = record["prop"]
                rel_type_props.setdefault(rel_type, []).append({
                    "name": prop,
                    "description": ""
                })

            # 构建 JSON
            json_obj = {
                "labels": [
                    {
                        "name": label,
                        "description": "",
                        "properties": label_props.get(label, [])
                    } for label in labels
                ],
                "relationships": [
                    {
                        "type": rel,
                        "description": "",
                        "properties": rel_type_props.get(rel, [])
                    } for rel in rel_types
                ],
                "triples": triples
            }

            # 保存到文件
            with open(output_path, "w", encoding="utf-8") as f:
                json.dump(json_obj, f, ensure_ascii=False, indent=2)

            return output_path

    #获取所有节点名称
    def get_all_node_names(self, label: str = None):
        """
        获取指定标签下所有节点的 name 属性
        :param label: 节点标签（如 'Effect', 'Symptom', 'Disease'）
        :return: List[str]
        """
        if label is None:
            query = """
            MATCH (n)
            RETURN DISTINCT n.name AS name
            ORDER BY name
            """
        else:
            query = f"""
            MATCH (n:{label})
            RETURN DISTINCT n.name AS name
            ORDER BY name
            """
        with self.driver.session() as session:
            result = session.run(query)
            return [record["name"] for record in result if record["name"]]

    def validate_cypher(self, query: str) -> bool:
        """
        检测 Cypher 查询语句是否合法（语法层面）
        :param query: 待检测的 Cypher 语句
        :return: True 表示合法，False 表示不合法
        """
        try:
            with self.driver.session() as session:
                # 使用 EXPLAIN 只做解析，不执行
                session.run(f"EXPLAIN {query}")
            return True
        except Exception as e:
            print(f"Cypher 语法错误: {e}")
            return False

neo4j_client = Neo4jClient(conf.NEO4J_URI, conf.NEO4J_USER, conf.NEO4J_PASSWORD)
if __name__ == '__main__':
    neo4j_client = Neo4jClient(conf.NEO4J_URI, conf.NEO4J_USER, conf.NEO4J_PASSWORD)
    # 查询所有节点 :cyper:MATCH (n) RETURN n
    # nodes = neo4j_client.run_cypher("MATCH (n) RETURN n.name")
    # print(nodes)

    # 测试一下 run_multiple_cypher
    # 注意 MATCH 要带 RETURN 才会有结果行；只写 MATCH 不返回任何东西，
    # 对应的结果项就是空列表（原来那条 MATCH (n:Employee ...) 就是这种情况）。
    # result = neo4j_client.run_multiple_cypher(
    #     [
    #         ("CREATE (n:Test {name: $name})", {"name": "Bob"}),
    #         ("MATCH (n:Test {name: $name}) RETURN n.name AS name", {"name": "Bob"}),
    #         ("MATCH (n:Test {name: $name}) DELETE n", {"name": "Bob"}),  # 自测收尾，别留脏数据
    #     ]
    # )
    # print(result)  # [[], [{'name': 'Bob'}], []]

    #测试export_tcm_metadata_to_json
    # neo4j_client.export_tcm_metadata_to_json()
    #测试 get_all_node_names
    # print(neo4j_client.get_all_node_names())
    #获取疾病下的所以node
    # print(neo4j_client.get_all_node_names("Disease"))
    result=neo4j_client.validate_cypher("MATCH1 (n:Disease) RETURN n.name")
    print(result)