"""
Neo4j Graph Database Helper Library
A comprehensive helper class for interacting with Neo4j graph database.
Provides connection management, query execution, graph operations, and vector search.
"""

import os
import asyncio
import json
import time
import logging
import hashlib
from datetime import datetime, timedelta
from typing import List, Dict, Optional, Union, Callable, Any, Set
from dataclasses import dataclass, field
from pathlib import Path

try:
    from neo4j import GraphDatabase, AsyncGraphDatabase
    from neo4j import (
        Driver,
        AsyncDriver,
        Session,
        AsyncSession,
        Transaction,
        AsyncTransaction,
        Result,
        AsyncResult,
        Record,
    )
    from neo4j.exceptions import (
        ServiceUnavailable,
        SessionExpired,
        AuthError,
        CypherSyntaxError,
        TransactionError,
    )
except ImportError:
    raise ImportError("Please install neo4j: pip install neo4j")

try:
    from cachetools import TTLCache
except ImportError:
    TTLCache = None


@dataclass
class Neo4jConfig:
    """Configuration class for Neo4j helper."""
    uri: str = "bolt://localhost:7687"
    username: str = "neo4j"
    password: str = "password"
    database: str = "neo4j"
    max_connection_lifetime: int = 3600
    max_connection_pool_size: int = 50
    connection_acquisition_timeout: int = 60
    max_retry_time: int = 30
    initial_retry_delay: float = 1.0
    retry_delay_multiplier: float = 2.0
    retry_delay_jitter_factor: float = 0.2
    enable_caching: bool = True
    cache_ttl: int = 3600
    log_level: str = "INFO"
    # Vector search settings
    vector_index_name: str = "vector_index"
    vector_dimension: int = 1536

    def __post_init__(self):
        # Default behavior: do not override here to ensure explicit values (e.g. from graph_config) are respected.
        # Call apply_env_overrides() explicitly if environment variable fallback is desired.
        pass

    def apply_env_overrides(self):
        """Apply environment variable overrides if current values are defaults."""
        if self.uri == "bolt://localhost:7687" and os.getenv("NEO4J_URI"):
            self.uri = os.getenv("NEO4J_URI")
        if self.username == "neo4j" and os.getenv("NEO4J_USERNAME"):
            self.username = os.getenv("NEO4J_USERNAME")
        if self.password == "password" and os.getenv("NEO4J_PASSWORD"):
            self.password = os.getenv("NEO4J_PASSWORD")
        if self.database == "neo4j" and os.getenv("NEO4J_DATABASE"):
            self.database = os.getenv("NEO4J_DATABASE")

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "Neo4jConfig":
        """
        Create a Neo4jConfig from a dictionary (e.g., a MongoDB graph_config document).
        Supports common field names used in graph configurations.
        """
        uri = data.get("uri") or data.get("host_url") or data.get("connection_string") or data.get("host") or data.get("url") or "bolt://localhost:7687"
        
        # Ensure scheme is present if it looks like just a host/IP
        if "://" not in uri:
            uri = f"bolt://{uri}"
            
        port = data.get("port")
        if port and f":{port}" not in uri:
            # Only append if port isn't already in the host part
            uri = f"{uri}:{port}"
        
        return cls(
            uri=uri,
            username=data.get("username") or data.get("user") or "neo4j",
            password=data.get("password") or data.get("pass") or "password",
            database=data.get("database") or data.get("db_name") or "neo4j",
            # Pass through other settings if present
            max_connection_lifetime=data.get("max_connection_lifetime", 3600),
            max_connection_pool_size=data.get("max_connection_pool_size", 50),
            connection_acquisition_timeout=data.get("connection_acquisition_timeout", 60),
        )


@dataclass
class CacheEntry:
    """Cache entry with TTL support."""
    data: Any
    created_at: datetime
    ttl: int

    def is_expired(self) -> bool:
        return datetime.now() > self.created_at + timedelta(seconds=self.ttl)


@dataclass
class UsageStats:
    """Usage statistics tracking for Neo4j operations."""
    total_queries: int = 0
    read_queries: int = 0
    write_queries: int = 0
    failed_queries: int = 0
    total_latency_ms: float = 0.0

    def record_query(self, is_write: bool = False, latency_ms: float = 0.0, success: bool = True):
        self.total_queries += 1
        if is_write:
            self.write_queries += 1
        else:
            self.read_queries += 1
        if not success:
            self.failed_queries += 1
        self.total_latency_ms += latency_ms

    def get_avg_latency(self) -> float:
        if self.total_queries == 0:
            return 0.0
        return self.total_latency_ms / self.total_queries


class Neo4jHelper:
    """Enhanced Neo4j helper with advanced features."""

    def __init__(self, config: Optional[Neo4jConfig] = None):
        self.config = config or Neo4jConfig()
        self._driver: Optional[Driver] = None
        self._async_driver: Optional[AsyncDriver] = None
        self._setup_logging()
        self._setup_cache()
        self._setup_usage_tracking()

    def _setup_logging(self):
        logging.basicConfig(
            level=getattr(logging, self.config.log_level.upper(), logging.INFO),
            format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
        )
        self.logger = logging.getLogger(self.__class__.__name__)

    def _setup_cache(self):
        if self.config.enable_caching and TTLCache:
            self._cache: Dict[str, CacheEntry] = {}
        else:
            self._cache = {}

    def _setup_usage_tracking(self):
        self.usage_stats = UsageStats()

    # ---------------- Connection Management ----------------
    def connect(self) -> "Neo4jHelper":
        """Establish connection to Neo4j."""
        if self._driver is None:
            uri = self.config.uri
            driver_kwargs = {
                "auth": (self.config.username, self.config.password),
                "max_connection_lifetime": self.config.max_connection_lifetime,
                "max_connection_pool_size": self.config.max_connection_pool_size,
                "connection_acquisition_timeout": self.config.connection_acquisition_timeout,
            }
            
            # Explicitly handle self-signed certificate schemes if driver fallback is needed
            if "+ssc" in uri:
                try:
                    import neo4j
                    # Neo4j 6.x uses TrustAll class and trusted_certificates keyword
                    if hasattr(neo4j, "TrustAll"):
                        uri = uri.replace("+ssc", "")
                        driver_kwargs["trusted_certificates"] = neo4j.TrustAll()
                        driver_kwargs["encrypted"] = True
                    # Neo4j 5.x uses TRUST_ALL_CERTIFICATES constant and trust keyword
                    elif hasattr(neo4j, "TRUST_ALL_CERTIFICATES"):
                        driver_kwargs["trust"] = neo4j.TRUST_ALL_CERTIFICATES
                        driver_kwargs["encrypted"] = True
                except (ImportError, AttributeError):
                    pass

            self._driver = GraphDatabase.driver(uri, **driver_kwargs)
            # Verify connectivity
            with self._driver.session(database=self.config.database) as session:
                session.run("RETURN 1")
            self.logger.info(f"Connected to Neo4j at {uri}")
        return self

    def connect_async(self) -> "Neo4jHelper":
        """Establish async connection to Neo4j."""
        if self._async_driver is None:
            uri = self.config.uri
            driver_kwargs = {
                "auth": (self.config.username, self.config.password),
                "max_connection_lifetime": self.config.max_connection_lifetime,
                "max_connection_pool_size": self.config.max_connection_pool_size,
                "connection_acquisition_timeout": self.config.connection_acquisition_timeout,
            }
            
            if "+ssc" in uri:
                try:
                    import neo4j
                    if hasattr(neo4j, "TrustAll"):
                        uri = uri.replace("+ssc", "")
                        driver_kwargs["trusted_certificates"] = neo4j.TrustAll()
                        driver_kwargs["encrypted"] = True
                    elif hasattr(neo4j, "TRUST_ALL_CERTIFICATES"):
                        driver_kwargs["trust"] = neo4j.TRUST_ALL_CERTIFICATES
                        driver_kwargs["encrypted"] = True
                except (ImportError, AttributeError):
                    pass

            self._async_driver = AsyncGraphDatabase.driver(uri, **driver_kwargs)
            self.logger.info(f"Async connected to Neo4j at {uri}")
        return self

    def close(self):
        """Close the Neo4j connection."""
        if self._driver:
            self._driver.close()
            self._driver = None
            self.logger.info("Neo4j connection closed")

    async def close_async(self):
        """Close the async Neo4j connection."""
        if self._async_driver:
            await self._async_driver.close()
            self._async_driver = None
            self.logger.info("Async Neo4j connection closed")

    def verify_connectivity(self) -> bool:
        """Verify that the connection is alive."""
        try:
            if self._driver is None:
                self.connect()
            with self._driver.session(database=self.config.database) as session:
                result = session.run("RETURN 1 AS num")
                result.single()
            return True
        except Exception as e:
            self.logger.error(f"Connectivity check failed: {e}")
            return False

    # ---------------- Cache helpers ----------------
    def _cache_key(self, *args, **kwargs) -> str:
        key_data = {"args": args, "kwargs": kwargs}
        key_str = json.dumps(key_data, sort_keys=True, default=str)
        return hashlib.md5(key_str.encode()).hexdigest()

    def _get_from_cache(self, key: str) -> Optional[Any]:
        if not self.config.enable_caching or not self._cache:
            return None
        entry = self._cache.get(key)
        if entry and not entry.is_expired():
            self.logger.debug(f"Cache hit for key: {key}")
            return entry.data
        if entry:
            self._cache.pop(key, None)
        return None

    def _set_cache(self, key: str, data: Any, ttl: Optional[int] = None):
        if not self.config.enable_caching or not self._cache:
            return
        ttl = ttl or self.config.cache_ttl
        self._cache[key] = CacheEntry(data=data, created_at=datetime.now(), ttl=ttl)
        self.logger.debug(f"Cache set for key: {key}")

    # ---------------- Query Execution ----------------
    def _execute_query(
        self,
        cypher: str,
        parameters: Optional[Dict[str, Any]] = None,
        write: bool = False,
        database: Optional[str] = None,
    ) -> List[Dict[str, Any]]:
        """Execute a Cypher query and return results as list of dicts."""
        if self._driver is None:
            self.connect()

        database = database or self.config.database
        parameters = parameters or {}
        start_time = time.time()

        try:
            with self._driver.session(database=database) as session:
                if write:
                    result = session.execute_write(self._run_tx, cypher, parameters)
                else:
                    result = session.execute_read(self._run_tx, cypher, parameters)

                latency_ms = (time.time() - start_time) * 1000
                self.usage_stats.record_query(is_write=write, latency_ms=latency_ms, success=True)
                return result

        except (ServiceUnavailable, SessionExpired) as e:
            self.logger.warning(f"Connection error, retrying: {e}")
            latency_ms = (time.time() - start_time) * 1000
            self.usage_stats.record_query(is_write=write, latency_ms=latency_ms, success=False)
            raise

        except Exception as e:
            self.logger.error(f"Query execution failed: {e}")
            latency_ms = (time.time() - start_time) * 1000
            self.usage_stats.record_query(is_write=write, latency_ms=latency_ms, success=False)
            raise

    def _run_tx(
        self,
        tx: Transaction,
        cypher: str,
        parameters: Dict[str, Any],
    ) -> List[Dict[str, Any]]:
        """Run a single query in a transaction."""
        result = tx.run(cypher, parameters)
        return [dict(record) for record in result]

    def execute_read(
        self,
        cypher: str,
        parameters: Optional[Dict[str, Any]] = None,
        database: Optional[str] = None,
        use_cache: bool = False,
    ) -> List[Dict[str, Any]]:
        """Execute a read query."""
        if use_cache and self.config.enable_caching:
            cache_key = self._cache_key("read", cypher, parameters)
            cached = self._get_from_cache(cache_key)
            if cached is not None:
                return cached
            result = self._execute_query(cypher, parameters, write=False, database=database)
            self._set_cache(cache_key, result)
            return result
        return self._execute_query(cypher, parameters, write=False, database=database)

    def execute_write(
        self,
        cypher: str,
        parameters: Optional[Dict[str, Any]] = None,
        database: Optional[str] = None,
    ) -> List[Dict[str, Any]]:
        """Execute a write query."""
        # Clear cache after write operations
        result = self._execute_query(cypher, parameters, write=True, database=database)
        self._cache.clear()
        return result

    def run_query(
        self,
        cypher: str,
        parameters: Optional[Dict[str, Any]] = None,
        database: Optional[str] = None,
    ) -> List[Dict[str, Any]]:
        """Execute an auto-commit query."""
        if self._driver is None:
            self.connect()

        database = database or self.config.database
        parameters = parameters or {}

        with self._driver.session(database=database) as session:
            result = session.run(cypher, parameters)
            return [dict(record) for record in result]

    # ---------------- Async Query Execution ----------------
    async def execute_read_async(
        self,
        cypher: str,
        parameters: Optional[Dict[str, Any]] = None,
        database: Optional[str] = None,
    ) -> List[Dict[str, Any]]:
        """Execute an async read query."""
        if self._async_driver is None:
            self.connect_async()

        database = database or self.config.database
        parameters = parameters or {}

        async with self._async_driver.session(database=database) as session:
            result = await session.execute_read(self._run_tx_async, cypher, parameters)
            return result

    async def execute_write_async(
        self,
        cypher: str,
        parameters: Optional[Dict[str, Any]] = None,
        database: Optional[str] = None,
    ) -> List[Dict[str, Any]]:
        """Execute an async write query."""
        if self._async_driver is None:
            self.connect_async()

        database = database or self.config.database
        parameters = parameters or {}

        async with self._async_driver.session(database=database) as session:
            result = await session.execute_write(self._run_tx_async, cypher, parameters)
            return result

    async def _run_tx_async(
        self,
        tx: AsyncTransaction,
        cypher: str,
        parameters: Dict[str, Any],
    ) -> List[Dict[str, Any]]:
        """Run a single query in an async transaction."""
        result = await tx.run(cypher, parameters)
        records = []
        async for record in result:
            records.append(dict(record))
        return records

    # ---------------- Node Operations ----------------
    def create_node(
        self,
        label: str,
        properties: Dict[str, Any],
        return_node: bool = True,
    ) -> Optional[Dict[str, Any]]:
        """Create a node with the given label and properties."""
        props_str = ", ".join([f"{k}: ${k}" for k in properties.keys()])
        cypher = f"CREATE (n:{label} {{{props_str}}})"
        if return_node:
            cypher += " RETURN n"
            result = self.execute_write(cypher, properties)
            return result[0].get("n") if result else None
        else:
            self.execute_write(cypher, properties)
            return None

    def merge_node(
        self,
        label: str,
        properties: Dict[str, Any],
        merge_on: Union[str, List[str]] = "id",
        return_node: bool = True,
    ) -> Optional[Dict[str, Any]]:
        """Merge a node based on one or more properties."""
        if isinstance(merge_on, str):
            merge_on = [merge_on]
            
        # Ensure all merge properties exist
        for prop in merge_on:
            if properties.get(prop) is None:
                return self.create_node(label, properties, return_node)
        
        merge_props_str = ", ".join([f"{p}: $merge_{p}" for p in merge_on])
        props_assignments = ", ".join([f"n.{k} = ${k}" for k in properties.keys()])
        cypher = f"MERGE (n:{label} {{{merge_props_str}}}) SET {props_assignments}"
        
        params = {**properties}
        for p in merge_on:
            params[f"merge_{p}"] = properties[p]
            
        if return_node:
            cypher += " RETURN n"
            result = self.execute_write(cypher, params)
            return result[0].get("n") if result else None
        else:
            self.execute_write(cypher, params)
            return None

    def get_node(
        self,
        label: str,
        properties: Dict[str, Any],
    ) -> Optional[Dict[str, Any]]:
        """Get a single node by label and properties."""
        props_str = " AND ".join([f"n.{k} = ${k}" for k in properties.keys()])
        cypher = f"MATCH (n:{label}) WHERE {props_str} RETURN n LIMIT 1"
        result = self.execute_read(cypher, properties)
        return result[0].get("n") if result else None

    def get_nodes(
        self,
        label: str,
        properties: Optional[Dict[str, Any]] = None,
        limit: int = 100,
    ) -> List[Dict[str, Any]]:
        """Get multiple nodes by label and optional properties."""
        props_str = ""
        params = properties or {}
        if properties:
            props_str = " AND ".join([f"n.{k} = ${k}" for k in properties.keys()])
            cypher = f"MATCH (n:{label}) WHERE {props_str} RETURN n LIMIT {limit}"
        else:
            cypher = f"MATCH (n:{label}) RETURN n LIMIT {limit}"
        result = self.execute_read(cypher, params)
        return [r.get("n", {}) for r in result]

    def update_node(
        self,
        label: str,
        match_properties: Dict[str, Any],
        update_properties: Dict[str, Any],
    ) -> int:
        """Update node properties and return number of updated nodes."""
        match_str = " AND ".join([f"n.{k} = ${k}" for k in match_properties.keys()])
        set_str = ", ".join([f"n.{k} = $update_{k}" for k in update_properties.keys()])
        cypher = f"MATCH (n:{label}) WHERE {match_str} SET {set_str}"
        params = {**match_properties, **{f"update_{k}": v for k, v in update_properties.items()}}
        result = self.execute_write(cypher, params)
        return len(result)

    def delete_node(
        self,
        label: str,
        properties: Dict[str, Any],
    ) -> int:
        """Delete nodes matching the given properties."""
        props_str = " AND ".join([f"n.{k} = ${k}" for k in properties.keys()])
        cypher = f"MATCH (n:{label}) WHERE {props_str} DETACH DELETE n"
        self.execute_write(cypher, properties)
        return 1

    def delete_all_nodes(self, label: str) -> int:
        """Delete all nodes with a given label."""
        cypher = f"MATCH (n:{label}) DETACH DELETE n"
        self.execute_write(cypher)
        return 1

    # ---------------- Relationship Operations ----------------
    def create_relationship(
        self,
        source_label: str,
        source_match: Dict[str, Any],
        target_label: str,
        target_match: Dict[str, Any],
        rel_type: str,
        rel_properties: Optional[Dict[str, Any]] = None,
        return_rel: bool = False,
    ) -> Optional[Dict[str, Any]]:
        """Create a relationship between two nodes."""
        source_match_str = " AND ".join([f"n.{k} = $src_{k}" for k in source_match.keys()])
        target_match_str = " AND ".join([f"m.{k} = $tgt_{k}" for k in target_match.keys()])

        rel_props_str = ""
        params = {**{f"src_{k}": v for k, v in source_match.items()},
                  **{f"tgt_{k}": v for k, v in target_match.items()}}

        if rel_properties:
            rel_props_str = " {" + ", ".join([f"{k}: $rel_{k}" for k in rel_properties.keys()]) + "}"
            params.update({f"rel_{k}": v for k, v in rel_properties.items()})

        cypher = f"""
        MATCH (n:{source_label}), (m:{target_label})
        WHERE {source_match_str} AND {target_match_str}
        CREATE (n)-[r:{rel_type}{rel_props_str}]->(m)
        """
        if return_rel:
            cypher += " RETURN r"
            result = self.execute_write(cypher, params)
            return result[0].get("r") if result else None
        else:
            self.execute_write(cypher, params)
            return None

    def merge_relationship(
        self,
        source_label: str,
        source_match: Dict[str, Any],
        target_label: str,
        target_match: Dict[str, Any],
        rel_type: str,
        rel_properties: Optional[Dict[str, Any]] = None,
    ) -> Optional[Dict[str, Any]]:
        """Merge a relationship between two nodes."""
        source_match_str = " AND ".join([f"n.{k} = $src_{k}" for k in source_match.keys()])
        target_match_str = " AND ".join([f"m.{k} = $tgt_{k}" for k in target_match.keys()])

        rel_props_str = ""
        params = {**{f"src_{k}": v for k, v in source_match.items()},
                  **{f"tgt_{k}": v for k, v in target_match.items()}}

        if rel_properties:
            rel_props_str = " {" + ", ".join([f"{k}: $rel_{k}" for k in rel_properties.keys()]) + "}"
            params.update({f"rel_{k}": v for k, v in rel_properties.items()})

        cypher = f"""
        MATCH (n:{source_label}), (m:{target_label})
        WHERE {source_match_str} AND {target_match_str}
        MERGE (n)-[r:{rel_type}{rel_props_str}]->(m)
        RETURN r
        """
        result = self.execute_write(cypher, params)
        return result[0].get("r") if result else None

    def get_relationships(
        self,
        source_label: str,
        source_match: Dict[str, Any],
        target_label: str,
        target_match: Dict[str, Any],
        rel_type: Optional[str] = None,
    ) -> List[Dict[str, Any]]:
        """Get relationships between two nodes."""
        source_match_str = " AND ".join([f"n.{k} = $src_{k}" for k in source_match.keys()])
        target_match_str = " AND ".join([f"m.{k} = $tgt_{k}" for k in target_match.keys()])

        rel_type_str = f":{rel_type}" if rel_type else ""

        params = {**{f"src_{k}": v for k, v in source_match.items()},
                  **{f"tgt_{k}": v for k, v in target_match.items()}}

        cypher = f"""
        MATCH (n:{source_label})-[r{rel_type_str}]->(m:{target_label})
        WHERE {source_match_str} AND {target_match_str}
        RETURN r
        """
        result = self.execute_read(cypher, params)
        return [r.get("r", {}) for r in result]

    def delete_relationship(
        self,
        source_label: str,
        source_match: Dict[str, Any],
        target_label: str,
        target_match: Dict[str, Any],
        rel_type: str,
    ) -> int:
        """Delete relationships between two nodes."""
        source_match_str = " AND ".join([f"n.{k} = $src_{k}" for k in source_match.keys()])
        target_match_str = " AND ".join([f"m.{k} = $tgt_{k}" for k in target_match.keys()])

        params = {**{f"src_{k}": v for k, v in source_match.items()},
                  **{f"tgt_{k}": v for k, v in target_match.items()}}

        cypher = f"""
        MATCH (n:{source_label})-[r:{rel_type}]->(m:{target_label})
        WHERE {source_match_str} AND {target_match_str}
        DELETE r
        """
        self.execute_write(cypher, params)
        return 1

    # ---------------- Path Operations ----------------
    def find_shortest_path(
        self,
        start_label: str,
        start_match: Dict[str, Any],
        end_label: str,
        end_match: Dict[str, Any],
        max_depth: int = 4,
    ) -> List[Dict[str, Any]]:
        """Find shortest path between two nodes."""
        start_match_str = " AND ".join([f"n.{k} = $start_{k}" for k in start_match.keys()])
        end_match_str = " AND ".join([f"m.{k} = $end_{k}" for k in end_match.keys()])

        params = {**{f"start_{k}": v for k, v in start_match.items()},
                  **{f"end_{k}": v for k, v in end_match.items()}}

        cypher = f"""
        MATCH (n:{start_label}), (m:{end_label})
        WHERE {start_match_str} AND {end_match_str}
        MATCH p = shortestPath((n)-[*1..{max_depth}]-(m))
        RETURN p
        """
        result = self.execute_read(cypher, params)
        return result

    def find_all_paths(
        self,
        start_label: str,
        start_match: Dict[str, Any],
        end_label: str,
        end_match: Dict[str, Any],
        max_depth: int = 4,
    ) -> List[Dict[str, Any]]:
        """Find all paths between two nodes."""
        start_match_str = " AND ".join([f"n.{k} = $start_{k}" for k in start_match.keys()])
        end_match_str = " AND ".join([f"m.{k} = $end_{k}" for k in end_match.keys()])

        params = {**{f"start_{k}": v for k, v in start_match.items()},
                  **{f"end_{k}": v for k, v in end_match.items()}}

        cypher = f"""
        MATCH (n:{start_label}), (m:{end_label})
        WHERE {start_match_str} AND {end_match_str}
        MATCH p = (n)-[*1..{max_depth}]-(m)
        RETURN p
        """
        result = self.execute_read(cypher, params)
        return result

    # ---------------- Vector Search ----------------
    def create_vector_index(
        self,
        index_name: str,
        label: str,
        property_name: str,
        dimension: int = 1536,
    ) -> bool:
        """Create a vector index for embeddings search."""
        cypher = f"""
        CALL db.index.vector.create(
            '{index_name}',
            '{label}.{property_name}',
            {dimension}
        )
        """
        try:
            self.execute_write(cypher)
            self.logger.info(f"Vector index '{index_name}' created")
            return True
        except Exception as e:
            self.logger.error(f"Failed to create vector index: {e}")
            return False

    def create_node_with_embedding(
        self,
        label: str,
        properties: Dict[str, Any],
        embedding_property: str = "embedding",
    ) -> Optional[Dict[str, Any]]:
        """Create a node with an embedding vector."""
        props_str = ", ".join([f"{k}: ${k}" for k in properties.keys()])
        cypher = f"""
        CREATE (n:{label} {{{props_str}}})
        RETURN n
        """
        result = self.execute_write(cypher, properties)
        return result[0].get("n") if result else None

    def search_similar_by_vector(
        self,
        index_name: str,
        embedding: List[float],
        top_k: int = 10,
        label: Optional[str] = None,
        property_filter: Optional[Dict[str, Any]] = None,
    ) -> List[Dict[str, Any]]:
        """Search for similar nodes using vector similarity."""
        cypher = f"""
        CALL db.index.vector.queryNodes('{index_name}', {top_k}, $embedding)
        YIELD node, score
        """

        params = {"embedding": embedding}

        if label:
            cypher = f"""
            CALL db.index.vector.queryNodes('{index_name}', {top_k}, $embedding)
            YIELD node, score
            WHERE any(label IN labels(node) WHERE label = '{label}')
            """

        if property_filter:
            filter_str = " AND ".join([f"node.{k} = $filter_{k}" for k in property_filter.keys()])
            cypher += f" AND {filter_str}"
            params.update({f"filter_{k}": v for k, v in property_filter.items()})

        cypher += " RETURN node, score"

        result = self.execute_read(cypher, params)
        return [{"node": r.get("node"), "score": r.get("score")} for r in result]

    # ---------------- Schema Operations ----------------
    def get_labels(self) -> List[str]:
        """Get all node labels in the database."""
        cypher = "CALL db.labels() YIELD label RETURN label"
        result = self.execute_read(cypher)
        return [r.get("label") for r in result]

    def get_relationship_types(self) -> List[str]:
        """Get all relationship types in the database."""
        cypher = "CALL db.relationshipTypes() YIELD relationshipType RETURN relationshipType"
        result = self.execute_read(cypher)
        return [r.get("relationshipType") for r in result]

    def get_property_keys(self) -> List[str]:
        """Get all property keys in the database."""
        cypher = "CALL db.propertyKeys() YIELD propertyKey RETURN propertyKey"
        result = self.execute_read(cypher)
        return [r.get("propertyKey") for r in result]

    def get_schema(self) -> Dict[str, Any]:
        """Get database schema information."""
        return {
            "labels": self.get_labels(),
            "relationship_types": self.get_relationship_types(),
            "property_keys": self.get_property_keys(),
        }

    # ---------------- Utility Methods ----------------
    def count_nodes(self, label: str) -> int:
        """Count nodes with a given label."""
        cypher = f"MATCH (n:{label}) RETURN count(n) AS count"
        result = self.execute_read(cypher)
        return result[0].get("count", 0) if result else 0

    def exists(self, label: str, properties: Dict[str, Any]) -> bool:
        """Check if a node exists."""
        props_str = " AND ".join([f"n.{k} = ${k}" for k in properties.keys()])
        cypher = f"MATCH (n:{label}) WHERE {props_str} RETURN count(n) > 0 AS exists"
        result = self.execute_read(cypher, properties)
        return result[0].get("exists", False) if result else False

    def get_degree(
        self,
        label: str,
        properties: Dict[str, Any],
        rel_type: Optional[str] = None,
        direction: str = "both",
    ) -> int:
        """Get the degree (number of relationships) of a node."""
        props_str = " AND ".join([f"n.{k} = ${k}" for k in properties.keys()])
        rel_str = ""
        if rel_type:
            if direction == "incoming":
                rel_str = f"<-[r:{rel_type}]-"
            elif direction == "outgoing":
                rel_str = f"-[r:{rel_type}]->"
            else:
                rel_str = f"-[r:{rel_type}]-"
        else:
            if direction == "incoming":
                rel_str = "<-[r]-"
            elif direction == "outgoing":
                rel_str = "-[r]->"
            else:
                rel_str = "-[r]-"

        cypher = f"MATCH (n:{label}){rel_str}(m) WHERE {props_str} RETURN count(r) AS degree"
        result = self.execute_read(cypher, properties)
        return result[0].get("degree", 0) if result else 0

    # ---------------- Introspection & utilities ----------------
    def get_usage_stats(self) -> Dict[str, Any]:
        return {
            "total_queries": self.usage_stats.total_queries,
            "read_queries": self.usage_stats.read_queries,
            "write_queries": self.usage_stats.write_queries,
            "failed_queries": self.usage_stats.failed_queries,
            "avg_latency_ms": round(self.usage_stats.get_avg_latency(), 2),
            "cache_size": len(self._cache),
        }

    def clear_cache(self):
        """Clear the cache."""
        self._cache.clear()
        self.logger.info("Cache cleared")

    def reset_usage_stats(self):
        """Reset usage statistics."""
        self.usage_stats = UsageStats()
        self.logger.info("Usage statistics reset")

    def __enter__(self):
        return self.connect()

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()
        stats = self.get_usage_stats()
        self.logger.info(f"Session complete. Usage: {stats}")


# Convenience functions
def create_node_params(**kwargs) -> Dict[str, Any]:
    """Create node properties dictionary."""
    return kwargs


def create_relationship_props(**kwargs) -> Dict[str, Any]:
    """Create relationship properties dictionary."""
    return kwargs
