"""Neo4j GDS adapter — benchmarks graph algorithms via Graph Data Science library."""

from __future__ import annotations

import time
from typing import Any

from benchmarks.config import NEO4J_PASSWORD, NEO4J_URI, NEO4J_USER


class Neo4jGdsAdapter:
    """Algorithm benchmark adapter using Neo4j GDS procedures.

    Two-phase approach matching our AGE+NetworkX pattern:
      Phase 1: gds.graph.project() → record projectMillis (≈ our to_networkx)
      Phase 2: Algorithm calls via GDS stream procedures
    """

    def __init__(self) -> None:
        self._driver: Any = None
        self._gds: Any = None
        self._graph_name = "bench"
        self._projected: Any = None  # GDS Graph object

    @property
    def name(self) -> str:
        return "neo4j_gds"

    def connect(self) -> None:
        from graphdatascience import GraphDataScience

        self._gds = GraphDataScience(NEO4J_URI, auth=(NEO4J_USER, NEO4J_PASSWORD))
        # Also keep a raw driver for data loading
        import neo4j

        self._driver = neo4j.GraphDatabase.driver(
            NEO4J_URI, auth=(NEO4J_USER, NEO4J_PASSWORD)
        )
        self._driver.verify_connectivity()

    def close(self) -> None:
        self._drop_projection()
        if self._gds:
            self._gds.close()
            self._gds = None
        if self._driver:
            self._driver.close()
            self._driver = None

    def _run(self, query: str, **params: Any) -> list[dict[str, Any]]:
        with self._driver.session() as session:
            result = session.run(query, **params)
            return [dict(record) for record in result]

    def _write(self, query: str, **params: Any) -> Any:
        with self._driver.session() as session:
            return session.execute_write(
                lambda tx: tx.run(query, **params).consume()
            )

    # -- Schema / data loading (reuse Neo4j adapter pattern) -----------------

    def setup_schema(self, vec_dims: int) -> None:
        self._write("MATCH (n) DETACH DELETE n")
        try:
            self._write(
                "CREATE INDEX node_bench_id IF NOT EXISTS "
                "FOR (n:Node) ON (n.bench_id)"
            )
        except Exception:
            pass

    def teardown(self) -> None:
        self._drop_projection()
        try:
            self._write("MATCH (n) DETACH DELETE n")
        except Exception:
            pass

    def load_nodes(self, nodes: list[dict[str, Any]]) -> int:
        batch_size = 1000
        count = 0
        for i in range(0, len(nodes), batch_size):
            batch = nodes[i : i + batch_size]
            self._write(
                "UNWIND $batch AS node "
                "CREATE (n:Node {"
                "  bench_id: node.bench_id,"
                "  category: node.category,"
                "  name: node.name,"
                "  score: node.score,"
                "  active: node.active"
                "})",
                batch=batch,
            )
            count += len(batch)
        return count

    def load_edges(self, edges: list[dict[str, Any]]) -> int:
        batch_size = 1000
        count = 0
        for i in range(0, len(edges), batch_size):
            batch = edges[i : i + batch_size]
            self._write(
                "UNWIND $batch AS edge "
                "MATCH (a:Node {bench_id: edge.source_id}), "
                "      (b:Node {bench_id: edge.target_id}) "
                "CREATE (a)-[:EDGE {"
                "  rel_type: edge.rel_type,"
                "  weight: edge.weight"
                "}]->(b)",
                batch=batch,
            )
            count += len(batch)
        return count

    # -- GDS graph projection -----------------------------------------------

    def _drop_projection(self) -> None:
        if self._projected is not None:
            try:
                self._projected.drop()
            except Exception:
                pass
            self._projected = None

    def project_graph(self) -> tuple[float, Any]:
        """Project the in-memory GDS graph. Returns (project_ms, graph_obj)."""
        self._drop_projection()

        start = time.perf_counter()
        G, result = self._gds.graph.project(
            self._graph_name,
            "Node",
            {"EDGE": {"orientation": "NATURAL", "properties": ["weight"]}},
        )
        elapsed_ms = (time.perf_counter() - start) * 1000.0

        self._projected = G
        return elapsed_ms, G

    def _ensure_projected(self) -> Any:
        if self._projected is None:
            _, G = self.project_graph()
            return G
        return self._projected

    # -- Centrality algorithms -----------------------------------------------

    def degree_centrality(self) -> list[dict[str, Any]]:
        G = self._ensure_projected()
        result = self._gds.degree.stream(G)
        return result.to_dict("records")

    def betweenness_centrality(self) -> list[dict[str, Any]]:
        G = self._ensure_projected()
        result = self._gds.betweenness.stream(G)
        return result.to_dict("records")

    def closeness_centrality(self) -> list[dict[str, Any]]:
        G = self._ensure_projected()
        # GDS uses "closeness" → beta.closeness or gds.closeness
        result = self._gds.closeness.stream(G)
        return result.to_dict("records")

    def eigenvector_centrality(self) -> list[dict[str, Any]]:
        G = self._ensure_projected()
        result = self._gds.eigenvector.stream(G)
        return result.to_dict("records")

    def harmonic_centrality(self) -> list[dict[str, Any]]:
        G = self._ensure_projected()
        result = self._gds.closeness.harmonic.stream(G)
        return result.to_dict("records")

    # -- Community algorithms ------------------------------------------------

    def louvain(self) -> list[dict[str, Any]]:
        G = self._ensure_projected()
        result = self._gds.louvain.stream(G)
        return result.to_dict("records")

    def label_propagation(self) -> list[dict[str, Any]]:
        G = self._ensure_projected()
        result = self._gds.labelPropagation.stream(G)
        return result.to_dict("records")

    def modularity_optimization(self) -> list[dict[str, Any]]:
        G = self._ensure_projected()
        result = self._gds.modularityOptimization.stream(G)
        return result.to_dict("records")

    def leiden(self) -> list[dict[str, Any]]:
        G = self._ensure_projected()
        result = self._gds.leiden.stream(G)
        return result.to_dict("records")

    # -- Path algorithms -----------------------------------------------------

    def shortest_paths(self, pairs: list[tuple[int, int]]) -> list[dict[str, Any]]:
        """Run Dijkstra shortest path for each pair (by bench_id)."""
        G = self._ensure_projected()
        results = []
        for src_bid, tgt_bid in pairs[:10]:
            try:
                # Resolve bench_id to Neo4j internal node IDs
                src_rows = self._run(
                    "MATCH (n:Node {bench_id: $bid}) RETURN id(n) AS nid",
                    bid=src_bid,
                )
                tgt_rows = self._run(
                    "MATCH (n:Node {bench_id: $bid}) RETURN id(n) AS nid",
                    bid=tgt_bid,
                )
                if not src_rows or not tgt_rows:
                    results.append({"source": src_bid, "target": tgt_bid, "cost": -1})
                    continue

                result = self._gds.shortestPath.dijkstra.stream(
                    G,
                    sourceNode=src_rows[0]["nid"],
                    targetNode=tgt_rows[0]["nid"],
                    relationshipWeightProperty="weight",
                )
                if len(result) > 0:
                    results.append({
                        "source": src_bid,
                        "target": tgt_bid,
                        "cost": float(result.iloc[0]["totalCost"]),
                    })
                else:
                    results.append({"source": src_bid, "target": tgt_bid, "cost": -1})
            except Exception:
                results.append({"source": src_bid, "target": tgt_bid, "cost": -1})
        return results

    # -- Ranking algorithms --------------------------------------------------

    def pagerank(self) -> list[dict[str, Any]]:
        G = self._ensure_projected()
        result = self._gds.pageRank.stream(G)
        return result.to_dict("records")

    # -- Similarity (link prediction functions) ------------------------------

    def jaccard(self, pairs: list[tuple[int, int]]) -> list[dict[str, Any]]:
        results = []
        for src_bid, tgt_bid in pairs[:10]:
            try:
                rows = self._run(
                    "MATCH (a:Node {bench_id: $src}), (b:Node {bench_id: $tgt}) "
                    "RETURN gds.similarity.jaccard("
                    "  [(a)-->(x) | id(x)], [(b)-->(x) | id(x)]"
                    ") AS score",
                    src=src_bid,
                    tgt=tgt_bid,
                )
                score = rows[0]["score"] if rows else 0.0
            except Exception:
                score = 0.0
            results.append({"u": src_bid, "v": tgt_bid, "score": score})
        return results

    def adamic_adar(self, pairs: list[tuple[int, int]]) -> list[dict[str, Any]]:
        """Adamic-Adar via Cypher (GDS doesn't have a direct stream for this)."""
        results = []
        for src_bid, tgt_bid in pairs[:10]:
            try:
                rows = self._run(
                    "MATCH (a:Node {bench_id: $src})--(c)--(b:Node {bench_id: $tgt}) "
                    "WITH a, b, c, size((c)--()) AS degree "
                    "WHERE degree > 1 "
                    "RETURN sum(1.0 / log(degree)) AS score",
                    src=src_bid,
                    tgt=tgt_bid,
                )
                score = rows[0]["score"] if rows and rows[0]["score"] else 0.0
            except Exception:
                score = 0.0
            results.append({"u": src_bid, "v": tgt_bid, "score": score})
        return results

    def common_neighbors(self, pairs: list[tuple[int, int]]) -> list[dict[str, Any]]:
        results = []
        for src_bid, tgt_bid in pairs[:10]:
            try:
                rows = self._run(
                    "MATCH (a:Node {bench_id: $src})-->(c)<--(b:Node {bench_id: $tgt}) "
                    "RETURN count(DISTINCT c) AS score",
                    src=src_bid,
                    tgt=tgt_bid,
                )
                score = rows[0]["score"] if rows else 0
            except Exception:
                score = 0
            results.append({"u": src_bid, "v": tgt_bid, "score": score})
        return results

    # -- Link prediction -----------------------------------------------------

    def preferential_attachment(self, pairs: list[tuple[int, int]]) -> list[dict[str, Any]]:
        results = []
        for src_bid, tgt_bid in pairs[:10]:
            try:
                rows = self._run(
                    "MATCH (a:Node {bench_id: $src}), (b:Node {bench_id: $tgt}) "
                    "RETURN size((a)--()) * size((b)--()) AS score",
                    src=src_bid,
                    tgt=tgt_bid,
                )
                score = rows[0]["score"] if rows else 0
            except Exception:
                score = 0
            results.append({"u": src_bid, "v": tgt_bid, "score": score})
        return results

    def resource_allocation(self, pairs: list[tuple[int, int]]) -> list[dict[str, Any]]:
        results = []
        for src_bid, tgt_bid in pairs[:10]:
            try:
                rows = self._run(
                    "MATCH (a:Node {bench_id: $src})-->(c)<--(b:Node {bench_id: $tgt}) "
                    "WITH c, size((c)--()) AS degree "
                    "WHERE degree > 0 "
                    "RETURN sum(1.0 / degree) AS score",
                    src=src_bid,
                    tgt=tgt_bid,
                )
                score = rows[0]["score"] if rows and rows[0]["score"] else 0.0
            except Exception:
                score = 0.0
            results.append({"u": src_bid, "v": tgt_bid, "score": score})
        return results

    # -- Embeddings ----------------------------------------------------------

    def node2vec(self) -> list[dict[str, Any]]:
        G = self._ensure_projected()
        result = self._gds.node2vec.stream(
            G,
            embeddingDimension=32,
            walkLength=30,
            walksPerNode=200,
        )
        return result.to_dict("records")
