"""Memgraph — графовый срез (узлы Atom + CodeSymbol, рёбра RELATES).

Memgraph совместим с Bolt-протоколом Neo4j — используем официальный neo4j драйвер.
Рёбра единообразны: (любой узел)-[:RELATES {kind, inferred, source}]->(любой узел).
Это позволяет логической инференции работать и над атомами, и над кодом.
"""
from __future__ import annotations
from neo4j import GraphDatabase
from . import config

_driver = None

def driver():
    global _driver
    if _driver is None:
        _driver = GraphDatabase.driver(
            f"bolt://{config.MEMGRAPH_HOST}:{config.MEMGRAPH_PORT}"
        )
    return _driver

def upsert_atom(atom_id: int, type: str, claim: str | None, status: str) -> None:
    with driver().session() as s:
        s.run(
            "MERGE (a:Atom {id: $id}) SET a.type = $type, a.claim = $claim, a.status = $status",
            id=atom_id, type=type, claim=claim or "", status=status,
        )

def upsert_link(src: int, dst: int, kind: str, *, inferred: bool = False,
                source: str | None = None, confidence: float = 0.5) -> None:
    with driver().session() as s:
        s.run(
            "MATCH (a:Atom {id: $src}), (b:Atom {id: $dst}) "
            "MERGE (a)-[r:RELATES {kind: $kind}]->(b) "
            "SET r.inferred = $inferred, r.source = $source, r.confidence = $confidence",
            src=src, dst=dst, kind=kind, inferred=inferred, source=source, confidence=confidence,
        )

def delete_atom(atom_id: int) -> None:
    with driver().session() as s:
        s.run("MATCH (a:Atom {id: $id}) DETACH DELETE a", id=atom_id)

def upsert_code_symbol(name: str, *, repo: str = "", path: str = "",
                       language: str = "", kind: str = "") -> None:
    """Узел CodeSymbol (дедуп по name+repo+path)."""
    with driver().session() as s:
        s.run(
            "MERGE (c:CodeSymbol {name: $name, repo: $repo, path: $path}) "
            "SET c.language = $language, c.kind = $kind",
            name=name, repo=repo, path=path, language=language, kind=kind,
        )

def upsert_code_dependency(src: str, dst: str, kind: str, *, repo: str = "", path: str = "") -> None:
    """Ребро зависимости кода: (src)-[:RELATES {kind}]->(dst)."""
    with driver().session() as s:
        # dst может быть внешним (builtin) — создаём как CodeSymbol без repo/path
        s.run(
            "MERGE (a:CodeSymbol {name: $src, repo: $repo, path: $path}) "
            "MERGE (b:CodeSymbol {name: $dst}) "
            "MERGE (a)-[r:RELATES {kind: $kind}]->(b) "
            "SET r.source = 'extraction'",
            src=src, dst=dst, kind=kind, repo=repo, path=path,
        )

def apply_rule_one_step(if_kind: str, then_kind: str) -> int:
    """Одношаговое правило: (a)-[:RELATES {kind: if}]->(b) ⇒ (a)-[:RELATES {kind: then}]->(b).

    Возвращает число созданных рёбер.
    """
    with driver().session() as s:
        result = s.run(
            "MATCH (a)-[r:RELATES {kind: $if_kind}]->(b) "
            "OPTIONAL MATCH (a)-[x:RELATES {kind: $then_kind}]->(b) "
            "WITH a, b, x WHERE x IS NULL "
            "CREATE (a)-[r2:RELATES {kind: $then_kind, inferred: true, source: 'inference'}]->(b) "
            "RETURN count(r2) AS n",
            if_kind=if_kind, then_kind=then_kind,
        )
        return result.single()["n"]

def apply_rule_chain(k1: str, k2: str, then_kind: str) -> int:
    """Композиция: (a)-[:RELATES {kind: k1}]->(b)-[:RELATES {kind: k2}]->(c)
    ⇒ (a)-[:RELATES {kind: then}]->(c). Транзитивность. Возвращает число созданных."""
    with driver().session() as s:
        result = s.run(
            "MATCH (a)-[:RELATES {kind: $k1}]->(b)-[:RELATES {kind: $k2}]->(c) "
            "OPTIONAL MATCH (a)-[x:RELATES {kind: $then_kind}]->(c) "
            "WITH a, b, c, x WHERE x IS NULL "
            "CREATE (a)-[r2:RELATES {kind: $then_kind, inferred: true, source: 'inference'}]->(c) "
            "RETURN count(r2) AS n",
            k1=k1, k2=k2, then_kind=then_kind,
        )
        return result.single()["n"]

def count() -> int:
    with driver().session() as s:
        return s.run("MATCH (a:Atom) RETURN count(a) AS n").single()["n"]

def count_edges() -> int:
    with driver().session() as s:
        return s.run("MATCH ()-[r:RELATES]->() RETURN count(r) AS n").single()["n"]

def count_code_symbols() -> int:
    with driver().session() as s:
        return s.run("MATCH (c:CodeSymbol) RETURN count(c) AS n").single()["n"]

def query_edges(kind: str | None = None, limit: int = 50) -> list[dict]:
    """Список рёбер (опционально по kind)."""
    with driver().session() as s:
        if kind:
            r = s.run(
                "MATCH (a)-[r:RELATES {kind: $kind}]->(b) "
                "RETURN a.name AS src, labels(a) AS al, b.name AS dst, labels(b) AS bl, "
                "r.kind AS kind, r.inferred AS inferred LIMIT $lim",
                kind=kind, lim=limit,
            )
        else:
            r = s.run(
                "MATCH (a)-[r:RELATES]->(b) "
                "RETURN a.name AS src, labels(a) AS al, b.name AS dst, labels(b) AS bl, "
                "r.kind AS kind, r.inferred AS inferred LIMIT $lim",
                lim=limit,
            )
        out = []
        for rec in r:
            d = {"src": rec["src"], "dst": rec["dst"]}
            if "al" in rec and rec["al"]:
                d["src_label"] = rec["al"][0]
            if "bl" in rec and rec["bl"]:
                d["dst_label"] = rec["bl"][0]
            for key in ("kind", "inferred"):
                try:
                    d[key] = rec[key]
                except KeyError:
                    pass
            out.append(d)
        return out