"""Kùzu adapter — embedded graph database via pip package."""

from __future__ import annotations

import concurrent.futures
import shutil
import time
from pathlib import Path
from typing import Any

import numpy as np

from benchmarks.config import KUZU_DB_PATH


class KuzuAdapter:
    """Benchmark adapter for Kùzu (embedded)."""

    def __init__(self) -> None:
        self._db: Any = None
        self._conn: Any = None
        self._db_path = KUZU_DB_PATH

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

    def connect(self) -> None:
        import kuzu

        # Clean start — Kuzu 0.11+ creates the path itself
        db_path = Path(self._db_path)
        if db_path.is_dir():
            shutil.rmtree(db_path)
        elif db_path.exists():
            db_path.unlink()

        self._db = kuzu.Database(str(db_path))
        self._conn = kuzu.Connection(self._db)

    def close(self) -> None:
        self._conn = None
        self._db = None

    def _execute(self, query: str, params: dict[str, Any] | None = None) -> list[dict[str, Any]]:
        """Execute a Cypher query and return results as dicts."""
        if params:
            result = self._conn.execute(query, params)
        else:
            result = self._conn.execute(query)
        rows = []
        while result.has_next():
            row = result.get_next()
            col_names = result.get_column_names()
            rows.append(dict(zip(col_names, row)))
        return rows

    def setup_schema(self, vec_dims: int) -> None:
        # Create node table
        self._conn.execute(
            "CREATE NODE TABLE Node("
            "  bench_id INT64 PRIMARY KEY,"
            "  category STRING,"
            "  name STRING,"
            "  score DOUBLE,"
            "  active BOOLEAN"
            ")"
        )

        # Create edge table
        self._conn.execute(
            "CREATE REL TABLE EDGE("
            "  FROM Node TO Node,"
            "  rel_type STRING,"
            "  weight DOUBLE,"
            "  timestamp INT64"
            ")"
        )

        # For write workloads
        self._conn.execute(
            "CREATE NODE TABLE NewNode("
            "  bench_id INT64 PRIMARY KEY,"
            "  name STRING"
            ")"
        )
        self._conn.execute(
            "CREATE REL TABLE NEW_EDGE("
            "  FROM Node TO Node,"
            "  weight DOUBLE"
            ")"
        )

    def teardown(self) -> None:
        self._conn = None
        self._db = None
        db_path = Path(self._db_path)
        if db_path.is_dir():
            shutil.rmtree(db_path)
        elif db_path.exists():
            db_path.unlink()

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

    def load_nodes(self, nodes: list[dict[str, Any]]) -> int:
        count = 0
        for node in nodes:
            self._conn.execute(
                "CREATE (n:Node {"
                "  bench_id: $bid,"
                "  category: $cat,"
                "  name: $name,"
                "  score: $score,"
                "  active: $active"
                "})",
                {
                    "bid": node["bench_id"],
                    "cat": node["category"],
                    "name": node["name"],
                    "score": node["score"],
                    "active": node["active"],
                },
            )
            count += 1
        return count

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

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

    _w01_counter: int = 0

    def w01_node_creation(self, nodes: list[dict[str, Any]]) -> int:
        count = 0
        for node in nodes:
            # Use unique IDs per iteration to avoid primary key collision
            uid = node["bench_id"] + (self._w01_counter * 1_000_000)
            self._conn.execute(
                "CREATE (n:NewNode {bench_id: $bid, name: $name})",
                {"bid": uid, "name": node["name"]},
            )
            count += 1
        self._w01_counter += 1
        return count

    def w02_edge_creation(self, edges: list[dict[str, Any]]) -> int:
        count = 0
        for edge in edges:
            try:
                self._conn.execute(
                    "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._execute(
            "MATCH (n:Node {bench_id: $bid}) "
            "RETURN n.bench_id, n.name, n.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._execute(
                "MATCH (a:Node {bench_id: $sid})-[:EDGE]->(b) RETURN b.bench_id",
                {"sid": start_id},
            )
        elif hops == 2:
            rows = self._execute(
                "MATCH (a:Node {bench_id: $sid})-[:EDGE]->(b)-[:EDGE]->(c) "
                "RETURN c.bench_id",
                {"sid": start_id},
            )
        else:
            rows = self._execute(
                "MATCH (a:Node {bench_id: $sid})-[:EDGE]->(b)-[:EDGE]->(c)-[:EDGE]->(d) "
                "RETURN d.bench_id",
                {"sid": start_id},
            )
        return len(rows)

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

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

    def w07_shortest_path(self, start_id: int, end_id: int) -> int:
        # Kuzu syntax: MATCH (a), (b), p = (a)-[:EDGE* SHORTEST 1..4]->(b)
        rows = self._execute(
            "MATCH (a:Node {bench_id: $src}), (b:Node {bench_id: $dst}), "
            "p = (a)-[:EDGE* SHORTEST 1..4]->(b) "
            "RETURN length(p)",
            {"src": start_id, "dst": end_id},
        )
        if rows:
            first_val = next(iter(rows[0].values()))
            return int(first_val)
        return -1

    def w08_pagerank(self, iterations: int) -> list[dict[str, Any]]:
        # Kùzu has built-in graph algorithms in some versions
        # Fall back to degree-based ranking
        rows = self._execute(
            "MATCH (n:Node)-[r:EDGE]-() "
            "RETURN n.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:
        # Kùzu does not have native vector search
        return 0

    def w10_hybrid_search(
        self, query_vec: np.ndarray, k: int, label: str
    ) -> int:
        # Not supported
        return 0

    _w11_counter: int = 0

    def w11_bulk_load(
        self, nodes: list[dict[str, Any]], edges: list[dict[str, Any]]
    ) -> int:
        # Offset node IDs to avoid primary key collisions across iterations
        offset = (self._w11_counter + 1) * 10_000_000
        offset_nodes = [
            {**n, "bench_id": n["bench_id"] + offset} for n in nodes
        ]
        n = self.load_nodes(offset_nodes)
        # Edges reference existing loaded nodes, not the offset ones
        e = self.load_edges(edges)
        self._w11_counter += 1
        return n + e

    def w12_concurrent(self, queries: list[str], concurrency: int) -> float:
        import kuzu

        start = time.perf_counter()

        def run_query(q: str) -> None:
            conn = kuzu.Connection(self._db)
            conn.execute(q)

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

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