"""Neo4j graph adapter (Phase 6)."""
from __future__ import annotations

from typing import Iterable, List

from ..types import GraphNode, GraphEdge


class Neo4jAdapter:
    """Minimal Neo4j adapter.

    Requires neo4j driver package.
    """

    def __init__(self, uri: str, user: str, password: str):
        try:
            from neo4j import GraphDatabase
        except Exception as exc:
            raise RuntimeError("neo4j driver not installed") from exc
        self.driver = GraphDatabase.driver(uri, auth=(user, password))

    def upsert_nodes(self, nodes: Iterable[GraphNode]) -> int:
        with self.driver.session() as session:
            for node in nodes:
                session.run(
                    "MERGE (n:KBNode {id: $id}) SET n.type = $type, n.properties = $props",
                    id=node.id, type=node.type, props=node.properties,
                )
        return len(list(nodes))

    def upsert_edges(self, edges: Iterable[GraphEdge]) -> int:
        with self.driver.session() as session:
            for edge in edges:
                session.run(
                    "MATCH (a:KBNode {id: $src}), (b:KBNode {id: $tgt}) "
                    "MERGE (a)-[r:KB_EDGE {type: $type}]->(b) "
                    "SET r.properties = $props",
                    src=edge.source, tgt=edge.target, type=edge.type, props=edge.properties,
                )
        return len(list(edges))

    def get_edges(self) -> List[GraphEdge]:
        edges: List[GraphEdge] = []
        with self.driver.session() as session:
            result = session.run("MATCH (a:KBNode)-[r:KB_EDGE]->(b:KBNode) RETURN a.id, b.id, r.type, r.properties")
            for row in result:
                edges.append(GraphEdge(source=row[0], target=row[1], type=row[2], properties=row[3] or {}))
        return edges
