"""NebulaGraph adapter — uses nebula3-python with hardcoded nGQL translations."""

from __future__ import annotations

import concurrent.futures
import time
from typing import Any

import numpy as np

from benchmarks.config import (
    NEBULA_HOST,
    NEBULA_PASSWORD,
    NEBULA_PORT,
    NEBULA_SPACE,
    NEBULA_USER,
)


class NebulaAdapter:
    """Benchmark adapter for NebulaGraph 3.x."""

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

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

    def connect(self) -> None:
        from nebula3.gclient.net import ConnectionPool
        from nebula3.Config import Config

        config = Config()
        config.max_connection_pool_size = 10
        self._pool = ConnectionPool()
        self._pool.init([(NEBULA_HOST, NEBULA_PORT)], config)
        self._session = self._pool.get_session(NEBULA_USER, NEBULA_PASSWORD)

        # Register storage hosts (required in NebulaGraph 3.x)
        try:
            self._session.execute('ADD HOSTS "nebula-storaged":9779')
        except Exception:
            pass  # May already be registered
        time.sleep(3)

    def close(self) -> None:
        if self._session:
            self._session.release()
            self._session = None
        if self._pool:
            self._pool.close()
            self._pool = None

    def _execute(self, stmt: str) -> Any:
        """Execute an nGQL statement."""
        result = self._session.execute(stmt)
        if not result.is_succeeded():
            raise RuntimeError(f"nGQL error: {result.error_msg()}")
        return result

    def _execute_rows(self, stmt: str) -> list[dict[str, Any]]:
        """Execute and return rows as dicts."""
        result = self._execute(stmt)
        rows = []
        if result.row_size() > 0:
            col_names = result.keys()
            for i in range(result.row_size()):
                row = {}
                for j, col in enumerate(col_names):
                    val = result.row_values(i)[j]
                    # Convert nebula Value types to Python
                    if val.is_int():
                        row[col] = val.as_int()
                    elif val.is_double():
                        row[col] = val.as_double()
                    elif val.is_string():
                        row[col] = val.as_string()
                    elif val.is_bool():
                        row[col] = val.as_bool()
                    else:
                        row[col] = str(val)
                rows.append(row)
        return rows

    def setup_schema(self, vec_dims: int) -> None:
        # Create space
        self._execute(
            f"CREATE SPACE IF NOT EXISTS {NEBULA_SPACE}"
            f"(vid_type=INT64, partition_num=10, replica_factor=1)"
        )
        # NebulaGraph needs time to propagate schema — retry USE
        for attempt in range(10):
            time.sleep(3)
            try:
                self._execute(f"USE {NEBULA_SPACE}")
                break
            except RuntimeError:
                if attempt == 9:
                    raise

        # Create tags (node types)
        self._execute(
            "CREATE TAG IF NOT EXISTS Node("
            "  bench_id int64,"
            "  category string,"
            "  name string,"
            "  score double,"
            "  active bool"
            ")"
        )
        self._execute(
            "CREATE TAG IF NOT EXISTS NewNode("
            "  bench_id int64,"
            "  name string"
            ")"
        )

        # Create edge types (ts instead of timestamp — reserved word in nGQL)
        self._execute(
            "CREATE EDGE IF NOT EXISTS EDGE_REL("
            "  rel_type string,"
            "  weight double,"
            "  ts int64"
            ")"
        )
        self._execute(
            "CREATE EDGE IF NOT EXISTS NEW_EDGE("
            "  weight double"
            ")"
        )

        # Create index
        time.sleep(6)
        self._execute(
            "CREATE TAG INDEX IF NOT EXISTS idx_node_bench_id ON Node(bench_id)"
        )
        self._execute(
            "CREATE TAG INDEX IF NOT EXISTS idx_node_category ON Node(category(32))"
        )
        time.sleep(6)

    def teardown(self) -> None:
        try:
            self._execute(f"DROP SPACE IF EXISTS {NEBULA_SPACE}")
        except Exception:
            pass

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

    def load_nodes(self, nodes: list[dict[str, Any]]) -> int:
        self._execute(f"USE {NEBULA_SPACE}")
        count = 0
        batch: list[str] = []
        for node in nodes:
            vid = node["bench_id"]
            active_str = "true" if node["active"] else "false"
            batch.append(
                f'{vid}:('
                f'{node["bench_id"]}, '
                f'"{node["category"]}", '
                f'"{node["name"]}", '
                f'{node["score"]}, '
                f'{active_str})'
            )
            count += 1
            if len(batch) >= 500:
                values = ", ".join(batch)
                self._execute(
                    f"INSERT VERTEX Node(bench_id, category, name, score, active) "
                    f"VALUES {values}"
                )
                batch = []
        if batch:
            values = ", ".join(batch)
            self._execute(
                f"INSERT VERTEX Node(bench_id, category, name, score, active) "
                f"VALUES {values}"
            )
        return count

    def load_edges(self, edges: list[dict[str, Any]]) -> int:
        self._execute(f"USE {NEBULA_SPACE}")
        count = 0
        batch: list[str] = []
        for edge in edges:
            batch.append(
                f'{edge["source_id"]}->{edge["target_id"]}:('
                f'"{edge["rel_type"]}", '
                f'{edge["weight"]}, '
                f'{edge["timestamp"]})'
            )
            count += 1
            if len(batch) >= 500:
                values = ", ".join(batch)
                self._execute(
                    f"INSERT EDGE EDGE_REL(rel_type, weight, ts) "
                    f"VALUES {values}"
                )
                batch = []
        if batch:
            values = ", ".join(batch)
            self._execute(
                f"INSERT EDGE EDGE_REL(rel_type, weight, ts) "
                f"VALUES {values}"
            )
        return count

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

    def w01_node_creation(self, nodes: list[dict[str, Any]]) -> int:
        self._execute(f"USE {NEBULA_SPACE}")
        count = 0
        for node in nodes:
            vid = node["bench_id"]
            self._execute(
                f'INSERT VERTEX NewNode(bench_id, name) '
                f'VALUES {vid}:({node["bench_id"]}, "{node["name"]}")'
            )
            count += 1
        return count

    def w02_edge_creation(self, edges: list[dict[str, Any]]) -> int:
        self._execute(f"USE {NEBULA_SPACE}")
        count = 0
        for edge in edges:
            try:
                self._execute(
                    f'INSERT EDGE NEW_EDGE(weight) '
                    f'VALUES {edge["source_id"]}->{edge["target_id"]}:({edge["weight"]})'
                )
                count += 1
            except Exception:
                pass
        return count

    def w03_point_lookup(self, node_id: int) -> dict[str, Any] | None:
        self._execute(f"USE {NEBULA_SPACE}")
        rows = self._execute_rows(
            f"FETCH PROP ON Node {node_id} "
            f"YIELD properties(vertex).bench_id AS bench_id, "
            f"properties(vertex).name AS name, "
            f"properties(vertex).score AS score"
        )
        return rows[0] if rows else None

    def w04_pattern_match(self, hops: int, start_id: int = 0) -> int:
        self._execute(f"USE {NEBULA_SPACE}")
        if hops == 1:
            rows = self._execute_rows(
                f"GO 1 STEPS FROM {start_id} OVER EDGE_REL "
                "YIELD dst(edge) AS bid"
            )
        elif hops == 2:
            rows = self._execute_rows(
                f"GO 1 TO 2 STEPS FROM {start_id} OVER EDGE_REL "
                "YIELD dst(edge) AS bid "
                "| GO 1 STEPS FROM $-.bid OVER EDGE_REL "
                "YIELD dst(edge) AS bid"
            )
        else:
            rows = self._execute_rows(
                f"GO 1 TO 3 STEPS FROM {start_id} OVER EDGE_REL "
                "YIELD dst(edge) AS bid"
            )
        return len(rows)

    def w05_aggregation(self) -> list[dict[str, Any]]:
        self._execute(f"USE {NEBULA_SPACE}")
        rows = self._execute_rows(
            "LOOKUP ON Node "
            "YIELD properties(vertex).category AS category "
            "| GROUP BY $-.category "
            "YIELD $-.category AS category, count(*) AS cnt"
        )
        return rows

    def w06_vle(self, min_hops: int, max_hops: int, start_id: int = 0) -> int:
        self._execute(f"USE {NEBULA_SPACE}")
        rows = self._execute_rows(
            f"GO {min_hops} TO {max_hops} STEPS FROM {start_id} OVER EDGE_REL "
            f"YIELD dst(edge) AS bid "
            f"| YIELD count(*) AS cnt"
        )
        return rows[0]["cnt"] if rows else 0

    def w07_shortest_path(self, start_id: int, end_id: int) -> int:
        self._execute(f"USE {NEBULA_SPACE}")
        rows = self._execute_rows(
            f"FIND SHORTEST PATH FROM {start_id} TO {end_id} "
            f"OVER EDGE_REL UPTO 4 STEPS "
            f"YIELD path AS p "
            f"| YIELD length($-.p) AS path_len"
        )
        return rows[0]["path_len"] if rows else -1

    def w08_pagerank(self, iterations: int) -> list[dict[str, Any]]:
        # NebulaGraph doesn't have built-in PageRank in Community
        # Fall back to degree count
        self._execute(f"USE {NEBULA_SPACE}")
        rows = self._execute_rows(
            "LOOKUP ON Node "
            "YIELD id(vertex) AS vid "
            "| GO 1 STEPS FROM $-.vid OVER EDGE_REL BIDIRECT "
            "YIELD $-.vid AS vid "
            "| GROUP BY $-.vid "
            "YIELD $-.vid AS bench_id, count(*) AS degree "
            "| ORDER BY $-.degree DESC | LIMIT 10"
        )
        return rows

    def w09_vector_search(self, query_vec: np.ndarray, k: int) -> int:
        # NebulaGraph does not support vector search
        return 0

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

    def w11_bulk_load(
        self, nodes: list[dict[str, Any]], edges: list[dict[str, Any]]
    ) -> int:
        n = self.load_nodes(nodes)
        e = self.load_edges(edges)
        return n + e

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

        def run_query(q: str) -> None:
            session = self._pool.get_session(NEBULA_USER, NEBULA_PASSWORD)
            try:
                session.execute(f"USE {NEBULA_SPACE}")
                # Translate simple Cypher to nGQL
                session.execute(
                    "FETCH PROP ON Node 1 "
                    "YIELD properties(vertex).bench_id AS bid"
                )
            finally:
                session.release()

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

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