src/graph_db/client.py

"""
Thin wrapper around Neo4j that exposes the same queries as test_db.py
but as Python methods usable by LangChain tools.
"""
from __future__ import annotations
from neo4j import GraphDatabase
from src.config import NEO4J_URI, NEO4J_USER, NEO4J_PASSWORD


class GraphClient:
    _instance: "GraphClient | None" = None

    def __init__(self):
        self.driver = self._new_driver()

    @staticmethod
    def _new_driver():
        return GraphDatabase.driver(
            NEO4J_URI,
            auth=(NEO4J_USER, NEO4J_PASSWORD),
            max_connection_lifetime=200,
        )

    @classmethod
    def get(cls) -> "GraphClient":
        if cls._instance is None:
            cls._instance = cls()
        return cls._instance

    def _run(self, cypher: str, **params) -> list[dict]:
        for attempt in range(2):
            try:
                with self.driver.session() as s:
                    return [dict(r) for r in s.run(cypher, **params)]
            except Exception as e:
                if attempt == 0 and "defunct" in str(e).lower():
                    self.driver.close()
                    self.driver = self._new_driver()
                else:
                    raise

    # ── Concept search ────────────────────────────────────────────────────────

    def search_concepts(self, keyword: str) -> list[dict]:
        rows = self._run("""
            CALL db.index.fulltext.queryNodes('concept_search', $kw)
            YIELD node, score
            OPTIONAL MATCH (node)-[:DEFINED_IN]->(s:Section)
            RETURN node.name       AS name,
                   node.definition AS definition,
                   node.example    AS example,
                   node.page       AS page,
                   s.id            AS section_id,
                   s.title         AS section_title,
                   score
            ORDER BY score DESC LIMIT 8
        """, kw=keyword)
        if not rows:
            # fallback: CONTAINS
            rows = self._run("""
                MATCH (c:Concept)
                WHERE toLower(c.name)       CONTAINS toLower($kw)
                   OR toLower(c.definition) CONTAINS toLower($kw)
                OPTIONAL MATCH (c)-[:DEFINED_IN]->(s:Section)
                RETURN c.name       AS name,
                       c.definition AS definition,
                       c.example    AS example,
                       c.page       AS page,
                       s.id         AS section_id,
                       s.title      AS section_title,
                       0.0          AS score
                LIMIT 8
            """, kw=keyword)
        return rows

    # ── Section queries ───────────────────────────────────────────────────────

    def chapter_info(self, chapter_id: str) -> dict | None:
        rows = self._run("""
            MATCH (ch:Chapter {id: $id})
            OPTIONAL MATCH (pt:Part)-[:CONTAINS]->(ch)
            OPTIONAL MATCH (ch)-[:CONTAINS]->(s:Section)
            RETURN ch.id       AS chapter_id,
                   ch.title    AS chapter_title,
                   pt.id       AS part_id,
                   pt.name     AS part_name,
                   s.id        AS section_id,
                   s.title     AS section_title,
                   s.page      AS page
            ORDER BY s.id
        """, id=chapter_id)
        if not rows:
            return None
        first = rows[0]
        sections = [
            {"section_id": r["section_id"], "title": r["section_title"], "page": r["page"]}
            for r in rows if r["section_id"]
        ]
        return {
            "chapter_id":    first["chapter_id"],
            "chapter_title": first["chapter_title"],
            "part_id":       first["part_id"],
            "part_name":     first["part_name"],
            "first_page":    sections[0]["page"] if sections else None,
            "sections":      sections,
        }

    def section_info(self, section_id: str) -> dict | None:
        rows = self._run("""
            MATCH (s:Section {id: $id})
            OPTIONAL MATCH (ch:Chapter)-[:CONTAINS]->(s)
            OPTIONAL MATCH (pt:Part)-[:CONTAINS]->(ch)
            RETURN s.id       AS section_id,
                   s.title    AS title,
                   s.page     AS page,
                   ch.id      AS chapter_id,
                   ch.title   AS chapter_title,
                   pt.id      AS part_id,
                   pt.name    AS part_name
        """, id=section_id)
        return rows[0] if rows else None

    def section_prereqs(self, section_id: str) -> list[dict]:
        return self._run("""
            MATCH path = (t:Section {id: $id})-[:REQUIRES*1..10]->(p:Section)
            WITH DISTINCT p, min(length(path)) AS depth
            RETURN p.id    AS section_id,
                   p.title AS title,
                   p.page  AS page,
                   depth
            ORDER BY depth, p.id
        """, id=section_id)

    def section_concepts(self, section_id: str) -> list[dict]:
        """Concepts defined IN a section."""
        return self._run("""
            MATCH (s:Section {id: $id})-[:DEFINES]->(c:Concept)
            RETURN c.name       AS name,
                   c.definition AS definition,
                   c.example    AS example,
                   c.page       AS page
            ORDER BY c.name
        """, id=section_id)

    def concepts_needed_for(self, section_id: str) -> list[dict]:
        """Prereq concepts for a section (via CONCEPT_REQUIRES, deduped, excl. same section)."""
        return self._run("""
            MATCH (s:Section {id: $id})-[:DEFINES]->(c:Concept)
            MATCH path = (c)-[:CONCEPT_REQUIRES*1..10]->(prereq:Concept)
            WHERE NOT (s)-[:DEFINES]->(prereq)
            WITH DISTINCT prereq, min(length(path)) AS depth
            OPTIONAL MATCH (prereq)-[:DEFINED_IN]->(ps:Section)
            RETURN prereq.name       AS name,
                   prereq.definition AS definition,
                   prereq.example    AS example,
                   prereq.page       AS page,
                   ps.id             AS section_id,
                   depth
            ORDER BY ps.id, prereq.name
        """, id=section_id)

    # ── Concept queries ───────────────────────────────────────────────────────

    def concept_prereqs(self, concept_name: str) -> list[dict]:
        return self._run("""
            MATCH (t:Concept)
            WHERE t.name = $name
               OR (t)-[:ALSO_KNOWN_AS]->(:Concept {name: $name})
            WITH t
            MATCH path = (t)-[:CONCEPT_REQUIRES*1..10]->(p:Concept)
            WITH DISTINCT p, min(length(path)) AS depth
            OPTIONAL MATCH (p)-[:DEFINED_IN]->(s:Section)
            RETURN p.name       AS name,
                   p.definition AS definition,
                   p.example    AS example,
                   p.page       AS page,
                   s.id         AS section_id,
                   s.title      AS section_title,
                   depth
            ORDER BY depth, s.id, p.name
        """, name=concept_name)

    def concept_subtypes(self, concept_name: str) -> list[dict]:
        return self._run("""
            MATCH (parent:Concept)
            WHERE parent.name = $name
               OR (parent)-[:ALSO_KNOWN_AS]->(:Concept {name: $name})
            WITH parent
            MATCH (child:Concept)-[:IS_TYPE_OF]->(parent)
            RETURN child.name       AS name,
                   child.definition AS definition,
                   child.example    AS example
            ORDER BY child.name
        """, name=concept_name)