MZC/KnowMathBot/backend/app/integrations/neo4j.py

"""
Neo4j 图数据库客户端：知识点/类别节点、CONTAINS 与 PREREQUISITE 关系的查询与维护。
"""
from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Optional

from neo4j import GraphDatabase
from neo4j.exceptions import Neo4jError

from app.core.config import settings


@dataclass(frozen=True)
class Neo4jClient:
    driver: Any

    def _session(self):
        return self.driver.session(database=settings.neo4j_database)

    def close(self) -> None:
        self.driver.close()

    def ensure_constraints(self) -> None:
        """创建唯一约束（KnowledgePoint.code、KnowledgeCategory.key），可重复调用。"""
        kp_cypher = """
        CREATE CONSTRAINT kp_code_unique IF NOT EXISTS
        FOR (k:KnowledgePoint) REQUIRE k.code IS UNIQUE
        """
        kc_cypher = """
        CREATE CONSTRAINT kc_key_unique IF NOT EXISTS
        FOR (c:KnowledgeCategory) REQUIRE c.key IS UNIQUE
        """
        with self._session() as session:
            session.run(kp_cypher)
            session.run(kc_cypher)

    def list_all_knowledge_points(self) -> list[dict[str, Any]]:
        """图中全部 KnowledgePoint，用于与前端目录对齐的下拉选项。"""
        cypher = """
        MATCH (k:KnowledgePoint)
        RETURN k.code AS code, k.name AS name, k.grade AS grade
        ORDER BY k.code
        """
        with self._session() as session:
            rows = session.run(cypher).data()
        out: list[dict[str, Any]] = []
        for row in rows:
            code = row.get("code")
            if not code:
                continue
            out.append(
                {
                    "code": str(code),
                    "name": row.get("name"),
                    "grade": row.get("grade"),
                }
            )
        return out

    def list_knowledge_catalog(self) -> dict[str, dict[str, list[dict[str, Any]]]]:
        """
        返回 Neo4j 中的知识点目录树：{grade: {category: [{code,name}, ...]}}
        与前端 KnowledgePage 的下拉层级一致。
        """
        cypher = """
        MATCH (c:KnowledgeCategory)-[:CONTAINS]->(k:KnowledgePoint)
        RETURN c.grade AS grade, c.name AS category, k.code AS code, k.name AS name
        ORDER BY grade, category, code
        """
        with self._session() as session:
            rows = session.run(cypher).data()

        out: dict[str, dict[str, list[dict[str, Any]]]] = {}
        for row in rows:
            grade = row.get("grade")
            category = row.get("category")
            code = row.get("code")
            name = row.get("name")
            if not grade or not category or not code:
                continue
            grade_s = str(grade)
            category_s = str(category)
            code_s = str(code)
            name_s = str(name) if name is not None else code_s
            out.setdefault(grade_s, {}).setdefault(category_s, []).append({"code": code_s, "name": name_s})
        return out

    def get_knowledge_point(self, code: str) -> Optional[dict[str, Any]]:
        cypher = """
        MATCH (k:KnowledgePoint {code: $code})
        OPTIONAL MATCH (c:KnowledgeCategory)-[:CONTAINS]->(k)
        RETURN k{ .code, .name, .definition, .grade, .subject } AS kp, c.name AS category
        """
        with self._session() as session:
            rec = session.run(cypher, code=code).single()
            if not rec:
                return None
            kp = dict(rec["kp"])
            kp["category"] = rec.get("category")
            return kp

    def list_sibling_knowledge_points(self, code: str, limit: int = 8) -> list[str]:
        """
        同类别下其它知识点（用于图预览 related），不含自身。
        返回格式：名称（编码）
        """
        cypher = """
        MATCH (k:KnowledgePoint {code: $code})
        OPTIONAL MATCH (c:KnowledgeCategory)-[:CONTAINS]->(k)
        WITH k, c
        WHERE c IS NOT NULL
        MATCH (c)-[:CONTAINS]->(other:KnowledgePoint)
        WHERE other.code <> k.code
        RETURN other.name AS name, other.code AS ocode
        ORDER BY ocode
        LIMIT $limit
        """
        with self._session() as session:
            rows = session.run(cypher, code=code, limit=limit).data()
        out: list[str] = []
        for row in rows:
            name = row.get("name") or row.get("ocode") or ""
            ocode = row.get("ocode") or ""
            if ocode:
                out.append(f"{name}（{ocode}）")
        return out

    def link_prerequisite(self, prerequisite_code: str, dependent_code: str) -> None:
        """
        (前置知识点)-[:PREREQUISITE]->(后置知识点)
        表示必须先掌握前置，再学后置。
        """
        pre = (prerequisite_code or "").strip()
        dep = (dependent_code or "").strip()
        if not pre or not dep:
            raise ValueError("前置与后置编码不能为空")
        if pre == dep:
            raise ValueError("前置与后置不能相同")
        if self.get_knowledge_point(pre) is None or self.get_knowledge_point(dep) is None:
            raise ValueError("前置或后置知识点在图中不存在，请先同步目录或写入知识点")
        cypher = """
        MATCH (pre:KnowledgePoint {code: $pre})
        MATCH (dep:KnowledgePoint {code: $dep})
        MERGE (pre)-[:PREREQUISITE]->(dep)
        """
        with self._session() as session:
            result = session.run(cypher, pre=pre, dep=dep)
            result.consume()

    def unlink_prerequisite(self, prerequisite_code: str, dependent_code: str) -> int:
        """删除一条 PREREQUISITE 边，返回删除条数（0 或 1）。"""
        pre = (prerequisite_code or "").strip()
        dep = (dependent_code or "").strip()
        if not pre or not dep:
            raise ValueError("前置与后置编码不能为空")
        cypher = """
        MATCH (pre:KnowledgePoint {code: $pre})-[r:PREREQUISITE]->(dep:KnowledgePoint {code: $dep})
        DELETE r
        RETURN 1 AS deleted
        """
        with self._session() as session:
            rec = session.run(cypher, pre=pre, dep=dep).single()
        return 1 if rec else 0

    def list_prerequisite_edges(self, code: str) -> tuple[list[str], list[str]]:
        """
        返回 (必须先掌握的前置列表, 以当前点为前置的后续列表)，展示格式：名称（编码）
        """
        code = (code or "").strip()
        if not code:
            return [], []

        inc_cypher = """
        MATCH (pre:KnowledgePoint)-[:PREREQUISITE]->(k:KnowledgePoint {code: $code})
        RETURN pre.name AS name, pre.code AS pcode
        ORDER BY pcode
        """
        out_cypher = """
        MATCH (k:KnowledgePoint {code: $code})-[:PREREQUISITE]->(nxt:KnowledgePoint)
        RETURN nxt.name AS name, nxt.code AS ncode
        ORDER BY ncode
        """
        incoming: list[str] = []
        outgoing: list[str] = []
        with self._session() as session:
            for row in session.run(inc_cypher, code=code).data():
                name = row.get("name") or row.get("pcode") or ""
                pcode = row.get("pcode") or ""
                if pcode:
                    incoming.append(f"{name}（{pcode}）")
            for row in session.run(out_cypher, code=code).data():
                name = row.get("name") or row.get("ncode") or ""
                ncode = row.get("ncode") or ""
                if ncode:
                    outgoing.append(f"{name}（{ncode}）")
        return incoming, outgoing

    def upsert_knowledge_point(self, code: str, name: str, grade: str | None = None, category: str | None = None) -> dict[str, Any]:
        if category:
            category_key = f"{grade or 'unknown'}::{category}"
            cypher = """
            MERGE (k:KnowledgePoint {code: $code})
            SET k.name = $name, k.grade = $grade
            MERGE (c:KnowledgeCategory {key: $category_key})
            SET c.name = $category, c.grade = $grade
            MERGE (c)-[:CONTAINS]->(k)
            RETURN k{ .code, .name, .definition, .grade, .subject } AS kp, c.name AS category
            """
            params = {"code": code, "name": name, "grade": grade, "category": category, "category_key": category_key}
        else:
            cypher = """
            MERGE (k:KnowledgePoint {code: $code})
            SET k.name = $name, k.grade = $grade
            RETURN k{ .code, .name, .definition, .grade, .subject } AS kp, null AS category
            """
            params = {"code": code, "name": name, "grade": grade}
        with self._session() as session:
            rec = session.run(cypher, **params).single()
            if not rec:
                raise RuntimeError("Neo4j upsert failed")
            kp = dict(rec["kp"])
            kp["category"] = rec.get("category")
            return kp

    def rename_category(self, grade: str, old_name: str, new_name: str) -> None:
        old_key = f"{grade}::{old_name}"
        new_key = f"{grade}::{new_name}"
        cypher = """
        MATCH (c:KnowledgeCategory {key: $old_key})
        SET c.name = $new_name, c.key = $new_key, c.grade = $grade
        RETURN c.key AS key
        """
        with self._session() as session:
            rec = session.run(cypher, old_key=old_key, new_key=new_key, new_name=new_name, grade=grade).single()
            if not rec:
                raise RuntimeError("Category not found")

    def delete_category(self, grade: str, name: str) -> None:
        key = f"{grade}::{name}"
        cypher = """
        MATCH (c:KnowledgeCategory {key: $key})-[:CONTAINS]->(k:KnowledgePoint)
        DETACH DELETE c, k
        """
        with self._session() as session:
            session.run(cypher, key=key)

    def rename_knowledge_point(self, code: str, new_name: str) -> None:
        cypher = """
        MATCH (k:KnowledgePoint {code: $code})
        SET k.name = $new_name
        RETURN k.code AS code
        """
        with self._session() as session:
            rec = session.run(cypher, code=code, new_name=new_name).single()
            if not rec:
                raise RuntimeError("Knowledge point not found")

    def delete_knowledge_point(self, code: str) -> None:
        cypher = """
        MATCH (k:KnowledgePoint {code: $code})
        DETACH DELETE k
        """
        with self._session() as session:
            session.run(cypher, code=code)

    def backfill_category_relations(self) -> dict[str, int]:
        fetch_cypher = """
        MATCH (k:KnowledgePoint)
        OPTIONAL MATCH (c:KnowledgeCategory)-[:CONTAINS]->(k)
        WITH k, c
        WHERE c IS NULL
        RETURN k.code AS code, k.name AS name
        """
        rule_map = {
            "J7-EQ": ("七年级", "代数与方程"),
            "J7-GEO": ("七年级", "几何基础"),
            "J8-FUNC": ("八年级", "代数与函数"),
            "J8-ROOT": ("八年级", "代数与函数"),
            "J8-GEO": ("八年级", "几何证明"),
            "J9-EQ": ("九年级", "代数进阶"),
            "J9-FUNC": ("九年级", "代数进阶"),
            "J9-CIRCLE": ("九年级", "圆与统计"),
            "J9-STAT": ("九年级", "圆与统计"),
        }

        def _infer(code: str) -> tuple[str, str]:
            for prefix, mapping in rule_map.items():
                if code.startswith(prefix):
                    return mapping
            if code.startswith("J7"):
                return ("七年级", "未分类")
            if code.startswith("J8"):
                return ("八年级", "未分类")
            if code.startswith("J9"):
                return ("九年级", "未分类")
            return ("未知年级", "未分类")

        with self._session() as session:
            rows = session.run(fetch_cypher).data()
            total = len(rows)
            updated = 0
            for row in rows:
                code = row.get("code")
                name = row.get("name") or code or "未命名知识点"
                if not code:
                    continue
                grade, category = _infer(code)
                self.upsert_knowledge_point(code=code, name=name, grade=grade, category=category)
                updated += 1
            return {"total": total, "updated": updated, "skipped": total - updated}


_client: Neo4jClient | None = None


def get_neo4j_client() -> Neo4jClient:
    global _client
    if _client is None:
        driver = GraphDatabase.driver(
            settings.neo4j_uri,
            auth=(settings.neo4j_user, settings.neo4j_password),
        )
        _client = Neo4jClient(driver=driver)
        try:
            _client.ensure_constraints()
        except Neo4jError:
            pass
    return _client