"""
Memgraph Graph Database Helper Library
A comprehensive helper class for interacting with Memgraph graph database.
Provides connection management, Cypher query execution, graph operations, and vector search.
Note: Memgraph is compatible with Neo4j's Cypher (OpenCypher), so the API is very similar.
"""

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:
    import mgclient
except ImportError:
    try:
        from gqlalchemy import Memgraph
    except ImportError:
        raise ImportError("Please install pymgclient or gqlalchemy: pip install pymgclient gqlalchemy")


@dataclass
class MemgraphConfig:
    """Configuration class for Memgraph helper."""
    host: str = "127.0.0.1"
    port: int = 7687
    username: str = ""
    password: str = ""
    database: str = ""
    ssl_mode: Optional[str] = None
    max_connection_lifetime: int = 3600
    max_connection_pool_size: int = 50
    connection_acquisition_timeout: int = 60
    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):
        # Env overrides
        self.host = os.getenv("MEMGRAPH_HOST", self.host)
        self.port = int(os.getenv("MEMGRAPH_PORT", str(self.port)))
        self.username = os.getenv("MEMGRAPH_USERNAME", self.username)
        self.password = os.getenv("MEMGRAPH_PASSWORD", self.password)
        self.database = os.getenv("MEMGRAPH_DATABASE", self.database)


@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 Memgraph 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 MemgraphHelper:
    """Enhanced Memgraph helper with advanced features."""

    def __init__(self, config: Optional[MemgraphConfig] = None):
        self.config = config or MemgraphConfig()
        self._connection = None
        self._cursor = 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):
        self._cache: Dict[str, CacheEntry] = {}

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

    # ---------------- Connection Management ----------------
    def connect(self) -> "MemgraphHelper":
        """Establish connection to Memgraph."""
        if self._connection is None:
            # Build connection string
            host = self.config.host
            port = self.config.port

            # SSL configuration
            sslmode = ""
            if self.config.ssl_mode:
                sslmode = f"?sslmode={self.config.ssl_mode}"

            connection_string = f"host={host} port={port}{sslmode}"

            try:
                # Try pymgclient first
                self._connection = mgclient.connect(
                    host=host,
                    port=port,
                    username=self.config.username if self.config.username else None,
                    password=self.config.password if self.config.password else None,
                    sslmode=self.config.ssl_mode,
                )
                self._connection.autocommit = True
                self._cursor = self._connection.cursor()
                self.logger.info(f"Connected to Memgraph at {host}:{port}")

            except (ImportError, Exception) as e:
                # Fallback: try gqlalchemy
                try:
                    from gqlalchemy import Memgraph as GQLMemgraph
                    self._memgraph = GQLMemgraph(
                        host=host,
                        port=port,
                        username=self.config.username,
                        password=self.config.password,
                        encrypted=self.config.ssl_mode is not None,
                    )
                    # Test connection
                    self._memgraph.execute("RETURN 1")
                    self.logger.info(f"Connected to Memgraph at {host}:{port} (via GQLAlchemy)")
                except Exception as e2:
                    raise ConnectionError(f"Failed to connect to Memgraph: {e2}")

        return self

    def connect_async(self) -> "MemgraphHelper":
        """Establish async connection to Memgraph."""
        # For pymgclient, we use the same connection
        return self.connect()

    def close(self):
        """Close the Memgraph connection."""
        if self._cursor:
            self._cursor.close()
            self._cursor = None
        if self._connection:
            self._connection.close()
            self._connection = None
        self.logger.info("Memgraph connection closed")

    async def close_async(self):
        """Close async connection."""
        self.close()

    def verify_connectivity(self) -> bool:
        """Verify that the connection is alive."""
        try:
            if self._connection is None:
                self.connect()
            self.execute("RETURN 1")
            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:
            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:
            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,
    ) -> List[Dict[str, Any]]:
        """Execute a Cypher query and return results as list of dicts."""
        if self._connection is None:
            self.connect()

        start_time = time.time()

        try:
            # Check if using gqlalchemy
            if hasattr(self, '_memgraph'):
                result = self._memgraph.execute(cypher, parameters or {})
                return self._format_gqlalchemy_result(result)
            else:
                # Use pymgclient
                if parameters:
                    # Convert parameters to proper format
                    for key, value in parameters.items():
                        if isinstance(value, dict):
                            parameters[key] = json.dumps(value)
                        elif not isinstance(value, (str, int, float, bool, list)):
                            parameters[key] = str(value)

                self._cursor.execute(cypher, parameters or {})

                # Fetch all results
                columns = [desc[0] for desc in self._cursor.description] if self._cursor.description else []
                rows = self._cursor.fetchall()

                results = []
                for row in rows:
                    row_dict = {}
                    for i, col in enumerate(columns):
                        value = row[i]
                        # Handle Memgraph types
                        if hasattr(value, 'value'):
                            value = value.value()
                        row_dict[col] = value
                    results.append(row_dict)

                latency_ms = (time.time() - start_time) * 1000
                is_write = cypher.strip().upper().startswith((
                    "INSERT", "UPDATE", "DELETE", "CREATE", "DROP", "MERGE", "SET"
                ))
                self.usage_stats.record_query(is_write=is_write, latency_ms=latency_ms, success=True)

                return results

        except Exception as e:
            latency_ms = (time.time() - start_time) * 1000
            is_write = cypher.strip().upper().startswith((
                "INSERT", "UPDATE", "DELETE", "CREATE", "DROP", "MERGE", "SET"
            ))
            self.usage_stats.record_query(is_write=is_write, latency_ms=latency_ms, success=False)
            self.logger.error(f"Query execution failed: {e}")
            raise

    def _format_gqlalchemy_result(self, result) -> List[Dict[str, Any]]:
        """Format GQLAlchemy result to list of dicts."""
        if hasattr(result, '__iter__'):
            return [dict(record) for record in result]
        return []

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

    def execute_read(
        self,
        cypher: str,
        parameters: Optional[Dict[str, Any]] = None,
        use_cache: bool = False,
    ) -> List[Dict[str, Any]]:
        """Execute a read query."""
        return self.execute(cypher, parameters, use_cache)

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

    # ---------------- 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: str = "id",
        return_node: bool = True,
    ) -> Optional[Dict[str, Any]]:
        """Merge a node based on a specific property."""
        merge_key = properties.get(merge_on, "")
        props_str = ", ".join([f"{k}: ${k}" for k in properties.keys()])

        if merge_key:
            cypher = f"MATCH (n:{label} {{{merge_on}: $merge_key}}) SET n = {{{props_str}}}"
        else:
            cypher = f"CREATE (n:{label} {{{props_str}}})"

        if return_node:
            cypher += " RETURN n"
            params = {**properties, "merge_key": merge_key}
            result = self.execute_write(cypher, params)
            return result[0].get("n") if result else None
        else:
            params = {**properties, "merge_key": merge_key}
            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."""
        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} = $match_{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 = {**{f"match_{k}": v for k, v in match_properties.items()},
                  **{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
        """
        return self.execute_read(cypher, params)

    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
        """
        return self.execute_read(cypher, params)

    # ---------------- Vector Search (Memgraph specific) ----------------
    def create_vector_index(
        self,
        index_name: str,
        label: str,
        property_name: str,
        dimension: int = 1536,
    ) -> bool:
        """Create a vector index for embeddings search."""
        # Note: Memgraph uses different vector search syntax
        cypher = f"""
        CREATE INDEX ON :{label}({property_name})
        """
        try:
            self.execute_write(cypher)
            self.logger.info(f"Index on '{label}({property_name})' created")
            return True
        except Exception as e:
            self.logger.error(f"Failed to create index: {e}")
            return False

    def search_similar_by_vector(
        self,
        label: str,
        embedding_property: str,
        embedding: List[float],
        top_k: int = 10,
        property_filter: Optional[Dict[str, Any]] = None,
    ) -> List[Dict[str, Any]]:
        """Search for similar nodes using vector similarity."""
        # Memgraph doesn't have native vector search in older versions
        # This is a placeholder - newer versions may have different syntax
        cypher = f"""
        MATCH (n:{label})
        RETURN n
        LIMIT {top_k}
        """
        return self.execute_read(cypher)

    # ---------------- Schema Operations ----------------
    def get_labels(self) -> List[str]:
        """Get all node labels in the database."""
        cypher = "CALL db.labels() YIELD label RETURN label"
        try:
            result = self.execute_read(cypher)
            return [r.get("label") for r in result]
        except:
            # Fallback for older Memgraph versions
            cypher = "MATCH (n) UNWIND labels(n) AS label RETURN DISTINCT 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"
        try:
            result = self.execute_read(cypher)
            return [r.get("relationshipType") for r in result]
        except:
            # Fallback
            cypher = "MATCH ()-[r]->() UNWIND type(r) AS relType RETURN DISTINCT relType"
            result = self.execute_read(cypher)
            return [r.get("relType") 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"
        try:
            result = self.execute_read(cypher)
            return [r.get("propertyKey") for r in result]
        except:
            return []

    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()])

        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

    # ---------------- Memgraph specific features ----------------
    def get_edge_properties(self) -> List[str]:
        """Get all edge properties (Memgraph specific)."""
        cypher = "CALL edge_property_keys() YIELD edge_property_key RETURN edge_property_key"
        try:
            result = self.execute_read(cypher)
            return [r.get("edge_property_key") for r in result]
        except:
            return []

    def clear_database(self) -> bool:
        """Clear all data in the database."""
        cypher = "MATCH (n) DETACH DELETE n"
        try:
            self.execute_write(cypher)
            self._cache.clear()
            return True
        except Exception as e:
            self.logger.error(f"Clear database failed: {e}")
            return False

    # ---------------- 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
