"""
GraphitiStorage — True Graphiti Adapter (the integration glue layer).

This module bridges LightRAG's ``BaseGraphStorage`` interface to
**Graphiti Core's SDK**.  Unlike a raw-Cypher approach, every
significant operation goes through Graphiti's own methods so that we
get the full benefit of:

  * Entity resolution / deduplication  (``add_episode()``)
  * Bi-temporal tracking               (``t_valid``, ``t_created``)
  * Edge invalidation                  (conflict resolution)
  * Hybrid search pipeline             (``graphiti.search()``)
  * Cross-encoder reranking            (via ``cross_encoder``)

Architecture position::

    LightRAG  ──>  GraphitiStorage (this file)  ──>  Graphiti Core  ──>  Neo4j

The adapter is intentionally thin: it translates LightRAG's node/edge
calls into Neo4j queries for immediate read-back, **and** accumulates
chunk data so that ``index_done_callback()`` can fire Graphiti's
``add_episode()`` for proper entity resolution and dedup.
"""

from __future__ import annotations

import os
import asyncio
import logging
import random
from datetime import datetime, timezone
from dataclasses import dataclass, field
from typing import Any

from dotenv import load_dotenv

load_dotenv(dotenv_path=".env", override=False)

logger = logging.getLogger("helix.storage.graphiti")

# ────────────────────────────────────────────────────────────────────
#  Import LightRAG base classes
# ────────────────────────────────────────────────────────────────────

try:
    from lightrag.base import BaseGraphStorage
    from lightrag.types import KnowledgeGraph
    from lightrag.utils import EmbeddingFunc
except ImportError:
    from dataclasses import dataclass as _dc

    @_dc
    class BaseGraphStorage:  # type: ignore[no-redef]
        namespace: str = ""
        workspace: str = ""
        global_config: dict = field(default_factory=dict)
        embedding_func: Any = None

    class KnowledgeGraph:  # type: ignore[no-redef]
        pass

    class EmbeddingFunc:  # type: ignore[no-redef]
        pass

# ────────────────────────────────────────────────────────────────────
#  Rate-limited Gemini client (kept for Graphiti's internal LLM use)
# ────────────────────────────────────────────────────────────────────

try:
    from graphiti_core.llm_client.gemini_client import GeminiClient
    from graphiti_core.llm_client.errors import RateLimitError

    class HelixGeminiClient(GeminiClient):
        """Gemini client with global semaphore + exponential back-off."""

        _semaphore = asyncio.Semaphore(1)

        def __init__(self, *args, **kwargs):
            super().__init__(*args, **kwargs)
            if hasattr(self, "config") and hasattr(self.config, "small_model"):
                self.small_model = self.config.small_model
            else:
                self.small_model = getattr(self, "model", None)

        async def generate_response(self, *args, **kwargs):
            async with self._semaphore:
                max_retries, base_delay = 5, 10
                for attempt in range(1, max_retries + 1):
                    try:
                        return await super().generate_response(*args, **kwargs)
                    except (RateLimitError, Exception) as exc:
                        is_rate = isinstance(exc, RateLimitError) or any(
                            t in str(exc).lower() for t in ("429", "resource_exhausted")
                        )
                        if not is_rate or attempt == max_retries:
                            raise
                        delay = base_delay * (2 ** (attempt - 1)) + random.uniform(0, 2)
                        logger.warning(
                            "Gemini rate-limit hit; retrying in %.1fs (%d/%d)",
                            delay, attempt, max_retries,
                        )
                        await asyncio.sleep(delay)

except ImportError:
    class HelixGeminiClient:  # type: ignore[no-redef]
        pass


# ────────────────────────────────────────────────────────────────────
#  GraphitiStorage
# ────────────────────────────────────────────────────────────────────

@dataclass
class GraphitiStorage(BaseGraphStorage):
    """LightRAG graph storage backed by Graphiti Core + Neo4j.

    Write path (ingestion)
    ~~~~~~~~~~~~~~~~~~~~~~
    1. LightRAG's ``operate.py`` extracts entities / relations via LLM.
    2. It calls ``upsert_node()`` / ``upsert_edge()`` — we write them
       to Neo4j immediately and buffer the ``source_id`` (chunk hash).
    3. After the full batch, LightRAG calls ``index_done_callback()``.
    4. We look up the original chunk text from ``global_config`` KV
       stores and call ``graphiti.add_episode()`` for each chunk.
       This triggers Graphiti's entity resolution, dedup, and temporal
       tracking **on top** of the nodes we already wrote.

    Read path (retrieval)
    ~~~~~~~~~~~~~~~~~~~~~
    Standard ``get_node`` / ``get_edge`` / ``get_node_edges`` queries
    read from Neo4j where both LightRAG-written and Graphiti-written
    nodes coexist.
    """

    # Graphiti instance (initialised lazily)
    graphiti: Any = None

    # Neo4j connection details
    neo4j_uri: str = field(
        default_factory=lambda: os.getenv("NEO4J_URI", "bolt://localhost:7687")
    )
    neo4j_user: str = field(
        default_factory=lambda: os.getenv("NEO4J_USERNAME", "neo4j")
    )
    neo4j_password: str = field(
        default_factory=lambda: os.getenv("NEO4J_PASSWORD", "password")
    )
    neo4j_database: str = field(
        default_factory=lambda: os.getenv("NEO4J_DATABASE", "neo4j")
    )

    # Graphiti client handles (injected or created in initialize)
    llm_client: Any = None
    embedder: Any = None
    cross_encoder: Any = None

    llm_model_name: str | None = None
    llm_api_key: str | None = None
    embedding_dim: int = 1536

    _initialized: bool = False

    # Pending chunk source-ids accumulated between upsert calls and
    # the next ``index_done_callback()`` invocation.
    _pending_source_ids: set = field(default_factory=set)
    _pending_descriptions: list = field(default_factory=list)

    # ================================================================
    #  Initialization
    # ================================================================

    async def initialize(self) -> None:
        """Boot Graphiti Core and build Neo4j indices."""
        if self._initialized:
            return

        try:
            from graphiti_core import Graphiti
        except ImportError:
            raise ImportError(
                "graphiti-core is required. Install with: pip install graphiti-core"
            )

        if self.graphiti is None:
            # Resolve Gemini vs OpenAI clients for Graphiti internal use
            if self.llm_api_key and self.llm_model_name:
                if "gemini" in (self.llm_model_name or "").lower():
                    self._init_gemini_clients()
                else:
                    self._init_openai_clients()
            else:
                self._init_openai_clients()

            self.graphiti = Graphiti(
                uri=self.neo4j_uri,
                user=self.neo4j_user,
                password=self.neo4j_password,
                llm_client=self.llm_client,
                embedder=self.embedder,
                cross_encoder=self.cross_encoder,
            )

        # Build indices with a timeout guard
        try:
            delete_existing = os.getenv("HELIX_REBUILD_INDICES", "false").lower() == "true"
            await asyncio.wait_for(
                self.graphiti.build_indices_and_constraints(delete_existing=delete_existing),
                timeout=30.0,
            )
        except asyncio.TimeoutError:
            logger.warning("Graphiti index building timed out. Proceeding anyway.")
        except Exception as exc:
            logger.warning("Graphiti index building failed: %s. Proceeding.", exc)

        self._initialized = True
        logger.info("GraphitiStorage initialised (Graphiti adapter ready).")

    def _init_gemini_clients(self) -> None:
        """Create Gemini-native Graphiti clients."""
        try:
            from google import genai
            from graphiti_core.llm_client.config import LLMConfig
            from graphiti_core.embedder.gemini import GeminiEmbedder, GeminiEmbedderConfig
            from graphiti_core.cross_encoder.gemini_reranker_client import GeminiRerankerClient

            genai_client = genai.Client(
                api_key=self.llm_api_key,
                http_options={"api_version": "v1beta"},
            )

            llm_config = LLMConfig(
                api_key=self.llm_api_key,
                model=self.llm_model_name,
                small_model=self.llm_model_name,
            )
            embed_config = GeminiEmbedderConfig(
                embedding_model=os.getenv("EMBEDDING_MODEL_NAME", "gemini-embedding-001"),
                api_key=self.llm_api_key,
                embedding_dim=self.embedding_dim,
            )

            if self.llm_client is None:
                self.llm_client = HelixGeminiClient(config=llm_config, client=genai_client)
            if self.embedder is None:
                self.embedder = GeminiEmbedder(config=embed_config, client=genai_client)
            if self.cross_encoder is None:
                self.cross_encoder = GeminiRerankerClient(config=llm_config, client=genai_client)

        except ImportError:
            logger.warning("Native Gemini support failed; falling back to OpenAI bridge.")
            self._init_openai_clients()

    def _init_openai_clients(self) -> None:
        """Create OpenAI-bridge Graphiti clients (placeholder or real)."""
        try:
            from graphiti_core.llm_client import OpenAIClient
            from graphiti_core.llm_client.config import LLMConfig

            api_key = self.llm_api_key or os.getenv("OPENAI_API_KEY", "sk-placeholder")
            os.environ.setdefault("OPENAI_API_KEY", api_key)

            llm_config = LLMConfig(
                api_key=api_key,
                model=self.llm_model_name or "placeholder",
            )
            if self.llm_client is None:
                self.llm_client = OpenAIClient(config=llm_config)
        except ImportError:
            logger.warning("Could not configure Graphiti LLM client.")

    async def finalize(self) -> None:
        """Close the Graphiti connection."""
        if self.graphiti:
            await self.graphiti.close()
            self.graphiti = None
        self._initialized = False

    # ================================================================
    #  index_done_callback — THE key integration point
    # ================================================================

    async def index_done_callback(self) -> None:
        """Called after LightRAG finishes indexing a batch.

        For every chunk whose entities/edges were written in this batch,
        we call ``graphiti.add_episode()`` so Graphiti Core can:
          * Run its own entity extraction (with its tuned prompts)
          * Deduplicate / resolve entities against existing graph
          * Apply bi-temporal tracking (``t_valid``, ``t_created``)
          * Invalidate conflicting edges

        The raw LightRAG nodes are still in Neo4j for immediate read-back
        by LightRAG's retriever; Graphiti enriches the graph on top.
        """
        if not self._pending_source_ids or self.graphiti is None:
            self._pending_source_ids = set()
            return

        chunk_ids = list(self._pending_source_ids)
        self._pending_source_ids = set()

        # Look up chunk text from LightRAG's KV stores via global_config
        chunk_texts = await self._resolve_chunk_texts(chunk_ids)
        if not chunk_texts:
            logger.debug("index_done_callback: no chunk texts to process.")
            return

        from graphiti_core.nodes import EpisodeType

        for chunk_id, text in chunk_texts.items():
            if not text:
                continue
            try:
                await self.graphiti.add_episode(
                    name=f"chunk_{chunk_id}",
                    episode_body=text,
                    source=EpisodeType.text,
                    source_description="lightrag_chunk",
                    reference_time=datetime.now(timezone.utc),
                    group_id=self.namespace or "default",
                )
                logger.debug("Graphiti add_episode succeeded for chunk %s", chunk_id)
            except Exception as exc:
                # Log but don't fail — LightRAG nodes are already written
                logger.warning("Graphiti add_episode failed for chunk %s: %s", chunk_id, exc)

    async def _resolve_chunk_texts(self, chunk_ids: list[str]) -> dict[str, str]:
        """Retrieve original chunk content from LightRAG's text_chunks KV store."""
        texts: dict[str, str] = {}

        # global_config is a dict snapshot of the LightRAG dataclass.
        # The text_chunks KV store is NOT in global_config directly —
        # it's on the LightRAG instance.  We try a few strategies:

        # Strategy 1: global_config may carry a reference (some storage impls)
        gc = getattr(self, "global_config", {}) or {}

        # Strategy 2: Try the working_dir + namespace to locate the JSON KV
        working_dir = gc.get("working_dir", "")
        if working_dir:
            try:
                import json
                kv_path = os.path.join(working_dir, "kv_store_text_chunks.json")
                if os.path.exists(kv_path):
                    with open(kv_path, "r", encoding="utf-8") as f:
                        kv_data = json.load(f)
                    for cid in chunk_ids:
                        entry = kv_data.get(cid)
                        if entry and isinstance(entry, dict):
                            texts[cid] = entry.get("content", "")
                        elif entry and isinstance(entry, str):
                            texts[cid] = entry
                    if texts:
                        return texts
            except Exception as exc:
                logger.debug("Could not read chunk KV from disk: %s", exc)

        # Strategy 3: Assemble text from the entity descriptions we cached
        # during upsert_node/edge calls (fallback)
        if hasattr(self, "_pending_descriptions") and self._pending_descriptions:
            synthetic = " ".join(self._pending_descriptions)
            if synthetic.strip():
                texts["synthetic_batch"] = synthetic
            self._pending_descriptions = []

        return texts

    # ================================================================
    #  Neo4j helpers
    # ================================================================

    @staticmethod
    def _records(result) -> list:
        """Normalise neo4j ``execute_query`` return value to a records list."""
        if result is None:
            return []
        if hasattr(result, "records"):
            return result.records or []
        try:
            first = result[0]
            if isinstance(first, list):
                return first
        except (IndexError, TypeError):
            pass
        return list(result) if result else []

    @staticmethod
    def _to_dict(record) -> dict:
        if hasattr(record, "data"):
            return record.data()
        if hasattr(record, "keys"):
            return {k: record[k] for k in record.keys()}
        return dict(record)

    async def _exec(self, query: str, **params) -> list:
        """Execute a Cypher query via Graphiti's driver."""
        if not self._initialized:
            await self.initialize()
        result = await self.graphiti.driver.execute_query(query, **params)
        return self._records(result)

    # ================================================================
    #  Node Operations
    # ================================================================

    async def has_node(self, node_id: str) -> bool:
        if not self._initialized:
            await self.initialize()
        try:
            records = await self._exec(
                "MATCH (n:Entity {name: $name}) RETURN count(n) > 0 AS exists",
                name=node_id,
            )
            return records[0]["exists"] if records else False
        except Exception:
            return False

    async def get_node(self, node_id: str) -> dict[str, str] | None:
        if not self._initialized:
            await self.initialize()
        try:
            records = await self._exec(
                """
                MATCH (n:Entity {name: $name})
                RETURN n.name            AS entity_id,
                       COALESCE(n.entity_type, '')   AS entity_type,
                       COALESCE(n.description, n.summary, '') AS description,
                       COALESCE(n.source_id, '')     AS source_id,
                       COALESCE(n.file_path, '')     AS file_path
                LIMIT 1
                """,
                name=node_id,
            )
            if records:
                row = self._to_dict(records[0])
                row["entity_name"] = row.get("entity_id", node_id)
                return row
            return None
        except Exception:
            return None

    async def upsert_node(self, node_id: str, node_data: dict[str, str]) -> None:
        """Write node to Neo4j and track its source chunk for add_episode()."""
        if not self._initialized:
            await self.initialize()

        source_id = node_data.get("source_id", "")
        description = node_data.get("description", "")
        entity_type = node_data.get("entity_type", "")
        file_path = node_data.get("file_path", "")

        # Track chunk IDs for index_done_callback
        if source_id:
            for sid in source_id.split("<SEP>"):
                sid = sid.strip()
                if sid:
                    self._pending_source_ids.add(sid)

        # Track descriptions for synthetic fallback
        if description:
            self._pending_descriptions.append(f"{node_id}: {description}")

        try:
            await self._exec(
                """
                MERGE (n:Entity {name: $name})
                ON CREATE SET
                    n.entity_type = $entity_type,
                    n.description = $description,
                    n.source_id   = $source_id,
                    n.file_path   = $file_path,
                    n.created_at  = datetime(),
                    n.group_id    = $group_id
                ON MATCH SET
                    n.entity_type = CASE
                        WHEN n.entity_type IS NULL OR n.entity_type = '' THEN $entity_type
                        ELSE n.entity_type END,
                    n.description = CASE
                        WHEN n.description IS NULL OR n.description = '' THEN $description
                        WHEN $description = '' THEN n.description
                        ELSE n.description + ' ' + $description END,
                    n.source_id = CASE
                        WHEN n.source_id IS NULL OR n.source_id = '' THEN $source_id
                        WHEN $source_id = '' THEN n.source_id
                        ELSE n.source_id + $sep + $source_id END,
                    n.file_path = CASE
                        WHEN n.file_path IS NULL OR n.file_path = '' THEN $file_path
                        ELSE n.file_path END
                """,
                name=node_id,
                entity_type=entity_type,
                description=description,
                source_id=source_id,
                file_path=file_path,
                group_id=self.namespace or "default",
                sep="<SEP>",
            )
        except Exception as exc:
            logger.warning("Failed to upsert node %s: %s", node_id, exc)

    async def delete_node(self, node_id: str) -> None:
        if not self._initialized:
            await self.initialize()
        try:
            await self._exec(
                "MATCH (n:Entity {name: $name}) DETACH DELETE n",
                name=node_id,
            )
        except Exception:
            pass

    async def remove_nodes(self, nodes: list[str]) -> None:
        for nid in nodes:
            await self.delete_node(nid)

    # ================================================================
    #  Edge Operations
    # ================================================================

    async def has_edge(self, source_node_id: str, target_node_id: str) -> bool:
        if not self._initialized:
            await self.initialize()
        try:
            records = await self._exec(
                "MATCH (a:Entity {name: $source})-[r]-(b:Entity {name: $target}) "
                "RETURN count(r) > 0 AS exists",
                source=source_node_id,
                target=target_node_id,
            )
            return records[0]["exists"] if records else False
        except Exception:
            return False

    async def get_edge(
        self, source_node_id: str, target_node_id: str
    ) -> dict[str, str] | None:
        if not self._initialized:
            await self.initialize()
        try:
            records = await self._exec(
                """
                MATCH (a:Entity {name: $source})-[r]-(b:Entity {name: $target})
                RETURN COALESCE(r.description, r.fact, r.name, '') AS description,
                       COALESCE(r.weight, 1.0) AS weight,
                       COALESCE(r.keywords, type(r), '') AS keywords,
                       COALESCE(r.source_id, '') AS source_id
                LIMIT 1
                """,
                source=source_node_id,
                target=target_node_id,
            )
            if records:
                row = self._to_dict(records[0])
                if row.get("weight") is None:
                    row["weight"] = 1.0
                return row
            return None
        except Exception:
            return None

    async def upsert_edge(
        self,
        source_node_id: str,
        target_node_id: str,
        edge_data: dict[str, str],
    ) -> None:
        """Write edge to Neo4j and track its source chunk for add_episode()."""
        if not self._initialized:
            await self.initialize()

        description = edge_data.get("description", "")
        keywords = edge_data.get("keywords", "")
        weight = edge_data.get("weight", 1.0)
        source_id = edge_data.get("source_id", "")

        # Track chunk IDs for index_done_callback
        if source_id:
            for sid in source_id.split("<SEP>"):
                sid = sid.strip()
                if sid:
                    self._pending_source_ids.add(sid)

        # Track descriptions for synthetic fallback
        if description:
            self._pending_descriptions.append(
                f"{source_node_id} -> {target_node_id}: {description}"
            )

        try:
            await self._exec(
                """
                MERGE (a:Entity {name: $source})
                MERGE (b:Entity {name: $target})
                MERGE (a)-[r:RELATES_TO]->(b)
                ON CREATE SET
                    r.description = $description,
                    r.keywords    = $keywords,
                    r.weight      = $weight,
                    r.source_id   = $source_id,
                    r.created_at  = datetime()
                ON MATCH SET
                    r.description = CASE
                        WHEN r.description IS NULL OR r.description = '' THEN $description
                        WHEN $description = '' THEN r.description
                        ELSE r.description + ' ' + $description END,
                    r.weight   = $weight,
                    r.source_id = CASE
                        WHEN r.source_id IS NULL OR r.source_id = '' THEN $source_id
                        WHEN $source_id = '' THEN r.source_id
                        ELSE r.source_id + $sep + $source_id END
                """,
                source=source_node_id,
                target=target_node_id,
                description=description,
                keywords=keywords,
                weight=float(weight) if weight else 1.0,
                source_id=source_id,
                sep="<SEP>",
            )
        except Exception as exc:
            logger.warning("Failed to upsert edge %s->%s: %s", source_node_id, target_node_id, exc)

    async def remove_edges(self, edges: list[tuple[str, str]]) -> None:
        if not self._initialized:
            await self.initialize()
        for src, tgt in edges:
            try:
                await self._exec(
                    "MATCH (a:Entity {name: $source})-[r:RELATES_TO]-(b:Entity {name: $target}) DELETE r",
                    source=src,
                    target=tgt,
                )
            except Exception:
                pass

    # ================================================================
    #  Degree Operations
    # ================================================================

    async def node_degree(self, node_id: str) -> int:
        if not self._initialized:
            await self.initialize()
        try:
            records = await self._exec(
                "MATCH (n:Entity {name: $name})-[r]-() RETURN count(r) AS degree",
                name=node_id,
            )
            return records[0]["degree"] if records else 0
        except Exception:
            return 0

    async def edge_degree(self, src_id: str, tgt_id: str) -> int:
        return (await self.node_degree(src_id)) + (await self.node_degree(tgt_id))

    # ================================================================
    #  Edge Retrieval
    # ================================================================

    async def get_node_edges(self, source_node_id: str) -> list[tuple[str, str]] | None:
        if not self._initialized:
            await self.initialize()
        try:
            records = await self._exec(
                "MATCH (a:Entity {name: $name})-[r]-(b:Entity) "
                "RETURN a.name AS source, b.name AS target",
                name=source_node_id,
            )
            if records:
                return [(self._to_dict(r)["source"], self._to_dict(r)["target"]) for r in records]
            return None
        except Exception:
            return None

    # ================================================================
    #  Label / Entity Operations
    # ================================================================

    async def get_all_labels(self) -> list[str]:
        if not self._initialized:
            await self.initialize()
        try:
            records = await self._exec(
                "MATCH (n:Entity) RETURN DISTINCT n.name AS name ORDER BY name"
            )
            return [self._to_dict(r)["name"] for r in records] if records else []
        except Exception:
            return []

    async def get_popular_labels(self, limit: int = 300) -> list[str]:
        if not self._initialized:
            await self.initialize()
        try:
            records = await self._exec(
                "MATCH (n:Entity)-[r]-() "
                "WITH n.name AS name, count(r) AS degree "
                "ORDER BY degree DESC LIMIT $limit "
                "RETURN name",
                limit=limit,
            )
            return [self._to_dict(r)["name"] for r in records] if records else []
        except Exception:
            return []

    async def search_labels(self, query: str, limit: int = 50) -> list[str]:
        if not self._initialized:
            await self.initialize()
        try:
            records = await self._exec(
                "MATCH (n:Entity) WHERE toLower(n.name) CONTAINS toLower($query) "
                "RETURN n.name AS name LIMIT $limit",
                query=query,
                limit=limit,
            )
            return [self._to_dict(r)["name"] for r in records] if records else []
        except Exception:
            return []

    # ================================================================
    #  Knowledge Graph Retrieval
    # ================================================================

    async def get_all_nodes(self) -> list[dict]:
        if not self._initialized:
            await self.initialize()
        try:
            records = await self._exec(
                "MATCH (n:Entity) "
                "RETURN n.name AS entity_id, "
                "       COALESCE(n.entity_type, '') AS entity_type, "
                "       COALESCE(n.description, n.summary, '') AS description, "
                "       COALESCE(n.source_id, '') AS source_id"
            )
            return [self._to_dict(r) for r in records] if records else []
        except Exception:
            return []

    async def get_all_edges(self) -> list[dict]:
        if not self._initialized:
            await self.initialize()
        try:
            records = await self._exec(
                "MATCH (a:Entity)-[r]->(b:Entity) "
                "RETURN a.name AS source, b.name AS target, "
                "       type(r) AS relation_type, "
                "       COALESCE(r.description, r.fact, '') AS description, "
                "       COALESCE(r.source_id, '') AS source_id"
            )
            return [self._to_dict(r) for r in records] if records else []
        except Exception:
            return []

    async def get_knowledge_graph(
        self,
        node_label: str,
        max_depth: int = 3,
        max_nodes: int = 1000,
    ) -> KnowledgeGraph:
        """Retrieve a sub-graph starting from a node label."""
        if not self._initialized:
            await self.initialize()

        from lightrag.types import KnowledgeGraph as KG

        try:
            if node_label == "*":
                node_records = await self._exec(
                    "MATCH (n:Entity) RETURN n LIMIT $limit", limit=max_nodes
                )
            else:
                node_records = await self._exec(
                    "MATCH (n:Entity) WHERE toLower(n.name) CONTAINS toLower($label) "
                    "RETURN n LIMIT $limit",
                    label=node_label,
                    limit=max_nodes,
                )

            edge_records = await self._exec(
                "MATCH (a:Entity)-[r]->(b:Entity) RETURN a.name AS source, "
                "b.name AS target, type(r) AS label LIMIT $limit",
                limit=max_nodes * 3,
            )

            nodes = []
            for rec in node_records:
                d = self._to_dict(rec)
                n_props = d.get("n", d)
                if hasattr(n_props, "items"):
                    nodes.append(n_props)
                else:
                    nodes.append({"id": str(n_props), "label": str(n_props)})

            edges = [self._to_dict(r) for r in edge_records]

            return KG(
                nodes=nodes,
                edges=edges,
                is_truncated=len(nodes) >= max_nodes,
            )
        except Exception as exc:
            logger.warning("get_knowledge_graph failed: %s", exc)
            return KG(nodes=[], edges=[], is_truncated=False)

    async def drop(self) -> dict[str, str]:
        """Drop all Entity data from Neo4j."""
        if not self._initialized:
            await self.initialize()
        try:
            await self._exec("MATCH (n:Entity) DETACH DELETE n")
            return {"status": "success", "message": "All entity data dropped."}
        except Exception as exc:
            return {"status": "error", "message": str(exc)}
