from __future__ import annotations

from typing import Any, Dict, Iterable, List


class Neo4jWriter:
    """
    Minimal Neo4j importer for ArchiMate graphs.

    This writer expects nodes/edges in the shape produced by `mapper.map_to_graph`.
    It intentionally avoids any dependency on the rest of the AIF project.
    """

    def __init__(self, uri: str, user: str, password: str):
        try:
            from neo4j import GraphDatabase  # type: ignore
        except ImportError as exc:  # pragma: no cover - import error path
            raise RuntimeError(
                "neo4j Python driver is required. "
                "Install it with `pip install neo4j`."
            ) from exc

        self._driver = GraphDatabase.driver(uri, auth=(user, password))

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

    def import_graph(
        self, nodes: Iterable[Dict[str, Any]], edges: Iterable[Dict[str, Any]]
    ) -> None:
        """
        Import nodes and edges into Neo4j.

        - Nodes are merged by `id` with label `ArchimateElement`.
        - Relationships are merged by `id` between existing nodes.
        """
        with self._driver.session() as session:
            # In small/medium models a simple loop is good enough.
            # For very large models, this can be batch-optimized later.
            for n in nodes:
                session.run(
                    """
                    MERGE (e:ArchimateElement {id: $id})
                    SET e.name = $name,
                        e.layer = $layer,
                        e += $properties
                    """,
                    id=n.get("id"),
                    name=n.get("name"),
                    layer=n.get("layer"),
                    properties=n.get("properties") or {},
                )

            for r in edges:
                session.run(
                    """
                    MATCH (s:ArchimateElement {id: $source})
                    MATCH (t:ArchimateElement {id: $target})
                    MERGE (s)-[rel:ARCHI_REL {id: $id}]->(t)
                    SET rel.type = $type,
                        rel += $properties
                    """,
                    id=r.get("id"),
                    type=r.get("type"),
                    source=r.get("source"),
                    target=r.get("target"),
                    properties=r.get("properties") or {},
                )

