"""Neo4j adapter — uses the official neo4j Python driver."""

from __future__ import annotations

import concurrent.futures
import time
from typing import Any

import numpy as np

from benchmarks.config import NEO4J_PASSWORD, NEO4J_URI, NEO4J_USER


class Neo4jAdapter:
    """Benchmark adapter for Neo4j 5.x Community."""

    def __init__(self) -> None:
        self._driver: Any = None

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

    def connect(self) -> None:
        import neo4j

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

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

    def _run(self, query: str, **params: Any) -> list[dict[str, Any]]:
        """Execute a Cypher query and return records as dicts."""
        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:
        """Execute a write transaction."""
        with self._driver.session() as session:
            return session.execute_write(
                lambda tx: tx.run(query, **params).consume()
            )

    def setup_schema(self, vec_dims: int) -> None:
        # Clear all data
        self._write("MATCH (n) DETACH DELETE n")

        # Create indexes
        try:
            self._write(
                "CREATE INDEX node_bench_id IF NOT EXISTS "
                "FOR (n:Node) ON (n.bench_id)"
            )
        except Exception:
            pass

        # Vector index (Neo4j 5.x supports native vector indexes)
        try:
            self._write(
                "CREATE VECTOR INDEX bench_vector IF NOT EXISTS "
                "FOR (n:Node) ON (n.embedding) "
                "OPTIONS {indexConfig: {"
                "  `vector.dimensions`: $dims,"
                "  `vector.similarity_function`: 'cosine'"
                "}}",
                dims=vec_dims,
            )
        except Exception:
            # Vector indexes may not be available in all Neo4j versions
            pass

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

    # -- Data loading ----------------------------------------------------

    def load_nodes(self, nodes: list[dict[str, Any]]) -> int:
        # Batch load with UNWIND
        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,"
                "  timestamp: edge.timestamp"
                "}]->(b)",
                batch=batch,
            )
            count += len(batch)
        return count

    # -- Workloads -------------------------------------------------------

    def w01_node_creation(self, nodes: list[dict[str, Any]]) -> int:
        count = 0
        for node in nodes:
            self._write(
                "CREATE (n:NewNode {bench_id: $bid, name: $name})",
                bid=node["bench_id"],
                name=node["name"],
            )
            count += 1
        return count

    def w02_edge_creation(self, edges: list[dict[str, Any]]) -> int:
        count = 0
        for edge in edges:
            try:
                self._write(
                    "MATCH (a:Node {bench_id: $src}), (b:Node {bench_id: $tgt}) "
                    "CREATE (a)-[:NEW_EDGE {weight: $w}]->(b)",
                    src=edge["source_id"],
                    tgt=edge["target_id"],
                    w=edge["weight"],
                )
                count += 1
            except Exception:
                pass
        return count

    def w03_point_lookup(self, node_id: int) -> dict[str, Any] | None:
        rows = self._run(
            "MATCH (n:Node {bench_id: $bid}) "
            "RETURN n.bench_id AS bench_id, n.name AS name, n.score AS score",
            bid=node_id,
        )
        return rows[0] if rows else None

    def w04_pattern_match(self, hops: int, start_id: int = 0) -> int:
        if hops == 1:
            rows = self._run(
                "MATCH (a:Node {bench_id: $sid})-[:EDGE]->(b) RETURN b.bench_id AS bid",
                sid=start_id,
            )
        elif hops == 2:
            rows = self._run(
                "MATCH (a:Node {bench_id: $sid})-[:EDGE]->(b)-[:EDGE]->(c) "
                "RETURN c.bench_id AS bid",
                sid=start_id,
            )
        else:
            rows = self._run(
                "MATCH (a:Node {bench_id: $sid})-[:EDGE]->(b)-[:EDGE]->(c)-[:EDGE]->(d) "
                "RETURN d.bench_id AS bid",
                sid=start_id,
            )
        return len(rows)

    def w05_aggregation(self) -> list[dict[str, Any]]:
        return self._run(
            "MATCH (n:Node) RETURN n.category AS category, count(*) AS cnt"
        )

    def w06_vle(self, min_hops: int, max_hops: int, start_id: int = 0) -> int:
        rows = self._run(
            f"MATCH (a:Node {{bench_id: $sid}})-[:EDGE*{min_hops}..{max_hops}]->(b) "
            "RETURN count(*) AS cnt",
            sid=start_id,
        )
        return rows[0]["cnt"] if rows else 0

    def w07_shortest_path(self, start_id: int, end_id: int) -> int:
        rows = self._run(
            "MATCH p = shortestPath("
            "(a:Node {bench_id: $start})-[:EDGE*..4]-(b:Node {bench_id: $end})"
            ") RETURN length(p) AS path_len",
            start=start_id,
            end=end_id,
        )
        return rows[0]["path_len"] if rows else -1

    def w08_pagerank(self, iterations: int) -> list[dict[str, Any]]:
        # Neo4j Community doesn't include GDS PageRank.
        # We approximate with a simple degree-based ranking.
        rows = self._run(
            "MATCH (n:Node)-[r]-() "
            "RETURN n.bench_id AS bench_id, count(r) AS degree "
            "ORDER BY degree DESC LIMIT 10"
        )
        return rows

    def w09_vector_search(self, query_vec: np.ndarray, k: int) -> int:
        try:
            rows = self._run(
                "CALL db.index.vector.queryNodes('bench_vector', $k, $vec) "
                "YIELD node, score "
                "RETURN node.bench_id AS bench_id, score",
                k=k,
                vec=query_vec.tolist(),
            )
            return len(rows)
        except Exception:
            # Vector index may not be configured
            return 0

    def w10_hybrid_search(
        self, query_vec: np.ndarray, k: int, label: str
    ) -> int:
        try:
            rows = self._run(
                "CALL db.index.vector.queryNodes('bench_vector', $k, $vec) "
                "YIELD node, score "
                "WHERE node.category = $label "
                "RETURN node.bench_id AS bench_id, score",
                k=k,
                vec=query_vec.tolist(),
                label=label,
            )
            return len(rows)
        except Exception:
            return 0

    def w11_bulk_load(
        self, nodes: list[dict[str, Any]], edges: list[dict[str, Any]]
    ) -> int:
        # Use UNWIND batch loading
        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=nodes,
        )
        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,"
            "  timestamp: edge.timestamp"
            "}]->(b)",
            batch=edges,
        )
        return len(nodes) + len(edges)

    def w12_concurrent(self, queries: list[str], concurrency: int) -> float:
        start = time.perf_counter()

        def run_query(q: str) -> None:
            with self._driver.session() as session:
                session.run(q).consume()

        with concurrent.futures.ThreadPoolExecutor(
            max_workers=concurrency
        ) as pool:
            list(pool.map(run_query, queries))

        elapsed = time.perf_counter() - start
        return len(queries) / elapsed
