"""把 _002_ 抽好的中药图谱 / 方剂图谱导入 Neo4j。

读 _002__extract_information/{中药图谱,方剂图谱}/*.json 里的 entities 和 relations，
在 Python 里拼 Cypher，走 common/neo4j_manager.py 的 Neo4jClient 写库。

几个关键取舍：

- **先在内存里去重，再写库。** 49226 条 entity 里只有 18870 个不同节点，50627 条
  relation 里只有 48346 条不同边 —— 同一个「咳嗽」在几百个文件里都出现。逐条 MERGE
  等于把同样的活干几百遍，所以先聚合成唯一集合再写。
- **节点主键是 (label, name)，不是 name。** 有 356 个名字同时挂在两种类型下
  （「热咳」既是 Disease 又是 Symptom），只按 name 去重会把两类实体粘成一个节点。
- **label 和关系类型只能拼进语句，不能当参数**，Cypher 不支持参数化它们。所以这两个
  位置都先对白名单做校验再拼字符串，值一律走 $ 参数，不做字符串插值。
- 属性以「实体自己那个文件」为准：一贯煎.json 里的「当归」只是配料，没有性味剂量，
  当归.json 里才是全的。同名属性冲突时后者覆盖前者，其余只补空缺。

用法（在项目根目录执行）：
    python -X utf8 -m _003_create_neo4j_database.01_import_graph_to_neo4j --limit 5
    python -X utf8 -m _003_create_neo4j_database.01_import_graph_to_neo4j --dry-run
    python -X utf8 -m _003_create_neo4j_database.01_import_graph_to_neo4j
    python -X utf8 -m _003_create_neo4j_database.01_import_graph_to_neo4j --only formula
"""

import argparse
import json
import os
from collections import Counter

from tqdm import tqdm

from common.config import Config
from common.neo4j_manager import Neo4jClient
from common.path_utils import get_file_path

# 白名单 = 唯一允许拼进 Cypher 的 label / 关系类型，和 _002_ 的 Pydantic schema 一致。
# 抽取结果里出现过 TREATS、TREATS_DISE 这种模型笔误，不在名单里就丢掉并记账。
ENTITY_TYPES = {"Symptom", "Disease", "Formula", "Herb", "Effect", "Source"}
RELATION_TYPES = {
    "TREATS_DISEASE",
    "ALLEVIATES_SYMPTOM",
    "HAS_EFFECT",
    "HAS_INGREDIENT",
    "HAS_SYMPTOM",
    "FROM_SOURCE",
}
# FormulaAttributes + HerbAttributes 的全部字段。属性名同样会拼进 SET，所以也过白名单，
# 免得模型多吐一个键就在库里长出一个野属性。
ATTR_KEYS = {
    "alias",
    "dosage",
    "effect",
    "indication",
    "meridian",
    "origin",
    "place",
    "processing",
    "property_flavor",
    "taboo",
    "traits",
    "usage",
}

GRAPH_DIRS = {
    "herb": "_002__extract_information/中药图谱",
    "formula": "_002__extract_information/方剂图谱",
}

# 一个事务里塞多少条语句。太小则往返次数多，太大则事务占内存、失败回滚的代价高。
BATCH_SIZE = 500


# ======================
# 读盘 + 聚合
# ======================
def iter_graph_files(only=None, limit=0):
    """按 (来源, 路径) 产出要导入的图谱 json 路径。"""
    paths = []
    for kind, rel_dir in GRAPH_DIRS.items():
        if only and kind != only:
            continue
        directory = get_file_path(rel_dir)
        if not os.path.isdir(directory):
            print(f"图谱目录不存在，跳过: {directory}")
            continue
        names = sorted(f for f in os.listdir(directory) if f.endswith(".json"))
        if limit:
            names = names[:limit]
        paths.extend((kind, os.path.join(directory, f)) for f in names)
    return paths


def clean_attributes(raw):
    """过滤属性字典：只留白名单里的键，丢掉空值，值统一成字符串。"""
    if not isinstance(raw, dict):
        return {}
    out = {}
    for key, value in raw.items():
        if key not in ATTR_KEYS or value is None:
            continue
        text = value if isinstance(value, str) else json.dumps(value, ensure_ascii=False)
        text = text.strip()
        if text:
            out[key] = text
    return out


def collect(paths, stats):
    """把所有图谱 json 聚合成「唯一节点」和「唯一关系」两个集合。

    nodes: {(label, name): {属性}}
    rels:  {(主label, 主name, 关系, 宾label, 宾name)}

    关系两端的节点也顺手补进 nodes —— 有 177 处 relation 的端点没在同文件的 entities
    里列出来，不补的话 MATCH 匹配不上，那条边就凭空丢了。
    """
    nodes = {}
    rels = set()

    for kind, path in tqdm(paths, desc="读取图谱", unit="个", ncols=90):
        try:
            with open(path, encoding="utf-8") as f:
                graph = json.load(f)
        except (OSError, json.JSONDecodeError):
            stats["json 读取失败"] += 1
            continue

        # 主体实体的属性最全（同名的配料条目往往只有名字），最后合并时优先它
        main_name = graph.get("name")

        for entity in graph.get("entities") or []:
            label = entity.get("type")
            name = (entity.get("name") or "").strip()
            if label not in ENTITY_TYPES or not name:
                stats["实体不合法丢弃"] += 1
                continue
            key = (label, name)
            attrs = clean_attributes(entity.get("attributes"))
            existing = nodes.setdefault(key, {})
            if name == main_name:
                # 自己那个文件里的定义最权威，直接盖掉别处带来的零碎属性
                existing.update(attrs)
            else:
                # 只补空缺，不覆盖已有的
                for k, v in attrs.items():
                    existing.setdefault(k, v)

        for relation in graph.get("relations") or []:
            rel_type = relation.get("relation")
            s_label, o_label = relation.get("subject_type"), relation.get("object_type")
            s_name = (relation.get("subject") or "").strip()
            o_name = (relation.get("object") or "").strip()
            if rel_type not in RELATION_TYPES:
                stats["关系类型不合法丢弃"] += 1
                continue
            if s_label not in ENTITY_TYPES or o_label not in ENTITY_TYPES:
                stats["关系端点类型不合法丢弃"] += 1
                continue
            if not s_name or not o_name:
                stats["关系端点为空丢弃"] += 1
                continue
            nodes.setdefault((s_label, s_name), {})
            nodes.setdefault((o_label, o_name), {})
            rels.add((s_label, s_name, rel_type, o_label, o_name))

    return nodes, rels


# ======================
# 拼 Cypher
# ======================
def build_constraint_cypher():
    """每类实体按 name 建唯一约束。

    顺带把 MERGE 需要的索引也建好了（约束自带索引），不然 18870 次 MERGE 都是全表扫。
    IF NOT EXISTS 让脚本可以重复跑。
    """
    return [
        (
            f"CREATE CONSTRAINT {label.lower()}_name IF NOT EXISTS "
            f"FOR (n:{label}) REQUIRE n.name IS UNIQUE",
            None,
        )
        for label in sorted(ENTITY_TYPES)
    ]


def build_node_cypher(nodes):
    """一个节点一条 MERGE。

    label 只能拼进语句（Cypher 不支持参数化 label），name 和属性全走 $ 参数。
    label 已经过了 ENTITY_TYPES 白名单，不存在注入面。

    属性用 SET n += $attrs 而不是 MERGE (n {name:..., effect:...})：属性写进 MERGE
    的话，属性一变就 MERGE 不上、变成新建一个节点。ON CREATE SET 也不合适 —— 重跑或
    补数据时已存在的节点就永远补不上属性了。
    """
    queries = []
    for (label, name), attrs in sorted(nodes.items()):
        if attrs:
            query = f"MERGE (n:{label} {{name: $name}}) SET n += $attrs"
            params = {"name": name, "attrs": attrs}
        else:
            query = f"MERGE (n:{label} {{name: $name}})"
            params = {"name": name}
        queries.append((query, params))
    return queries


def build_relation_cypher(rels):
    """一条关系一条语句：两端 MATCH（节点已经在上一步 MERGE 好了）再 MERGE 边。

    关系类型和两端 label 都拼进语句，同样先过了白名单。MERGE 而不是 CREATE，
    保证重复跑不会长出平行边。

    两个 MATCH 中间夹一个 WITH，而不是写成 `MATCH (s), (o)`：后者是一条语句里两个不
    相连的模式，Neo4j 每条都会回一个 cartesian product 警告，5 万条刷得日志没法看。
    这里两端都是唯一约束上的等值查找，笛卡尔积本来就是 1×1，换个写法只是让驱动闭嘴。
    """
    queries = []
    for s_label, s_name, rel_type, o_label, o_name in sorted(rels):
        query = (
            f"MATCH (s:{s_label} {{name: $s_name}}) "
            f"WITH s MATCH (o:{o_label} {{name: $o_name}}) "
            f"MERGE (s)-[:{rel_type}]->(o)"
        )
        queries.append((query, {"s_name": s_name, "o_name": o_name}))
    return queries


# ======================
# 写库
# ======================
def run_in_batches(client, queries, desc):
    """分批提交。

    run_multiple_cypher 整批是一个事务，5 万条一次性提交内存吃不消、失败也全回滚，
    所以切成 BATCH_SIZE 一批。批与批之间不是原子的，但所有语句都是 MERGE，
    中断后重跑是幂等的，无所谓。
    """
    total = len(queries)
    done = 0
    for start in range(0, total, BATCH_SIZE):
        batch = queries[start : start + BATCH_SIZE]
        client.run_multiple_cypher(batch)
        done += len(batch)
        print(f"  {desc}: {done}/{total}")
    return done


def run(only=None, limit=0, dry_run=False, wipe=False):
    paths = iter_graph_files(only=only, limit=limit)
    if not paths:
        print("没有找到图谱 json，先跑 _002_ 的抽取脚本。")
        return

    stats = Counter()
    nodes, rels = collect(paths, stats)

    label_count = Counter(label for label, _ in nodes)
    rel_count = Counter(r[2] for r in rels)

    print(f"\n图谱文件 {len(paths)} 个")
    print(f"唯一节点 {len(nodes)} 个: {dict(label_count.most_common())}")
    print(f"唯一关系 {len(rels)} 条: {dict(rel_count.most_common())}")
    for key, value in stats.most_common():
        print(f"  {key}: {value}")

    node_queries = build_node_cypher(nodes)
    rel_queries = build_relation_cypher(rels)

    if dry_run:
        print("\n--dry-run，不写库。示例语句：")
        for query, params in node_queries[:2] + rel_queries[:2]:
            print(f"  {query}\n    {params}")
        return

    conf = Config()
    client = Neo4jClient(conf.NEO4J_URI, conf.NEO4J_USER, conf.NEO4J_PASSWORD)

    if wipe:
        # 只删这个图谱的六类节点，不动库里别的东西。DETACH 连边一起删。
        print("\n清空旧数据...")
        for label in sorted(ENTITY_TYPES):
            client.run_cypher(f"MATCH (n:{label}) DETACH DELETE n")

    print("\n建唯一约束...")
    client.run_multiple_cypher(build_constraint_cypher())

    print("写入节点...")
    run_in_batches(client, node_queries, "节点")

    print("写入关系...")
    run_in_batches(client, rel_queries, "关系")

    # 从库里回读一遍计数，别只信自己发了多少条语句
    print("\n库内实际计数：")
    for label in sorted(ENTITY_TYPES):
        rows = client.run_cypher(f"MATCH (n:{label}) RETURN count(n) AS c")
        print(f"  {label}: {rows[0]['c']}")
    for rel_type in sorted(RELATION_TYPES):
        rows = client.run_cypher(f"MATCH ()-[r:{rel_type}]->() RETURN count(r) AS c")
        print(f"  {rel_type}: {rows[0]['c']}")


def main():
    parser = argparse.ArgumentParser(description="把 _002_ 抽取的图谱导入 Neo4j")
    parser.add_argument(
        "--only",
        choices=sorted(GRAPH_DIRS),
        help="只导入一边：herb=中药图谱，formula=方剂图谱，默认两边都导",
    )
    parser.add_argument(
        "--limit",
        type=int,
        default=0,
        help="每个目录只取前 N 个 json 试跑，默认 0 表示全量",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="只统计并打印示例 Cypher，不连库不写入",
    )
    parser.add_argument(
        "--wipe",
        action="store_true",
        help="导入前先 DETACH DELETE 这六类节点（会删掉已有图谱数据）",
    )
    args = parser.parse_args()
    run(only=args.only, limit=args.limit, dry_run=args.dry_run, wipe=args.wipe)


if __name__ == "__main__":
    main()