"""
Amazon Neptune Graph Database Helper Library
A comprehensive helper class for interacting with Amazon Neptune graph database.
Supports both Gremlin (property graph) and SPARQL (RDF) queries.
"""

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 boto3
    from botocore.exceptions import ClientError
except ImportError:
    raise ImportError("Please install boto3: pip install boto3")

try:
    from gremlin_python.driver import client, serializer
    from gremlin_python.driver.driver_remote_connection import DriverRemoteConnection
    from gremlin_python.process.anonymous_traversal import traversal
    from gremlin_python.process.graph_traversal import __
    from gremlin_python.process.traversal import T
except ImportError:
    gremlin = None


@dataclass
class NeptuneConfig:
    """Configuration class for Amazon Neptune helper."""
    # Connection settings
    host: str = "localhost"
    port: int = 8182
    region: str = "us-east-1"
    use_ssl: bool = True

    # Authentication
    iam_enabled: bool = False
    aws_access_key_id: Optional[str] = None
    aws_secret_access_key: Optional[str] = None
    aws_session_token: Optional[str] = None

    # Credentials
    username: str = ""
    password: str = ""

    # Database
    database: str = "graph"  # Neptune DB cluster identifier

    # Gremlin client settings
    gremlin_message_serializer: str = "graphsonv3"
    pool_size: int = 10

    # Cache settings
    enable_caching: bool = True
    cache_ttl: int = 3600
    log_level: str = "INFO"

    # Query mode
    query_mode: str = "gremlin"  # gremlin, sparql, or opencypher

    def __post_init__(self):
        # Env overrides
        self.host = os.getenv("NEPTUNE_HOST", self.host)
        self.port = int(os.getenv("NEPTUNE_PORT", str(self.port)))
        self.region = os.getenv("AWS_REGION", self.region)
        self.iam_enabled = os.getenv("NEPTUNE_IAM_ENABLED", "false").lower() == "true"
        self.aws_access_key_id = os.getenv("AWS_ACCESS_KEY_ID", self.aws_access_key_id)
        self.aws_secret_access_key = os.getenv("AWS_SECRET_ACCESS_KEY", self.aws_secret_access_key)
        self.aws_session_token = os.getenv("AWS_SESSION_TOKEN", self.aws_session_token)
        self.username = os.getenv("NEPTUNE_USERNAME", self.username)
        self.password = os.getenv("NEPTUNE_PASSWORD", self.password)
        self.database = os.getenv("NEPTUNE_DATABASE", self.database)
        self.use_ssl = os.getenv("NEPTUNE_USE_SSL", "true").lower() == "true"


@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 Neptune 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 NeptuneHelper:
    """Enhanced Amazon Neptune helper with advanced features."""

    def __init__(self, config: Optional[NeptuneConfig] = None):
        self.config = config or NeptuneConfig()
        self._gremlin_client = None
        self._neptune_data_client = 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) -> "NeptuneHelper":
        """Establish connection to Amazon Neptune."""
        if self.config.query_mode == "gremlin":
            self._connect_gremlin()
        return self

    def _connect_gremlin(self):
        """Connect to Neptune using Gremlin."""
        if gremlin is None:
            raise ImportError("Please install gremlinpython: pip install gremlinpython")

        if self._gremlin_client is None:
            # Build connection URL
            scheme = "wss" if self.config.use_ssl else "ws"
            url = f"{scheme}://{self.config.host}:{self.config.port}/gremlin"

            # Connection configuration
            if self.config.username and self.config.password:
                # Basic authentication
                self._gremlin_client = client.Client(
                    url,
                    "g",
                    message_serializer=serializer.GraphSONSerializersV3d0(),
                    username=self.config.username,
                    password=self.config.password,
                    pool_size=self.config.pool_size,
                )
            else:
                # No authentication
                self._gremlin_client = client.Client(
                    url,
                    "g",
                    message_serializer=serializer.GraphSONSerializersV3d0(),
                    pool_size=self.config.pool_size,
                )

            self.logger.info(f"Connected to Neptune at {self.config.host}:{self.config.port}")

    def _get_neptune_client(self):
        """Get or create Neptune Data API client."""
        if self._neptune_data_client is None:
            kwargs = {"region_name": self.config.region}

            if self.config.iam_enabled:
                # Use IAM authentication
                pass  # boto3 handles IAM automatically
            elif self.config.aws_access_key_id:
                kwargs["aws_access_key_id"] = self.config.aws_access_key_id
                kwargs["aws_secret_access_key"] = self.config.aws_secret_access_key
                if self.config.aws_session_token:
                    kwargs["aws_session_token"] = self.config.aws_session_token

            self._neptune_data_client = boto3.client("neptunedata", **kwargs)

        return self._neptune_data_client

    def connect_async(self) -> "NeptuneHelper":
        """Establish async connection to Neptune."""
        return self.connect()

    def close(self):
        """Close the Neptune connection."""
        if self._gremlin_client:
            self._gremlin_client.close()
            self._gremlin_client = None
            self.logger.info("Neptune Gremlin 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.config.query_mode == "gremlin":
                if self._gremlin_client is None:
                    self.connect()
                result = self._gremlin_client.submit("g.V().count()").all().result()
                return True
            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}")

    # ---------------- Gremlin Query Execution ----------------
    def execute_gremlin(
        self,
        query: str,
        parameters: Optional[Dict[str, Any]] = None,
        use_cache: bool = False,
    ) -> List[Any]:
        """Execute a Gremlin query."""
        if self._gremlin_client is None:
            self._connect_gremlin()

        start_time = time.time()

        try:
            if parameters:
                # Convert parameters to Gremlin format
                result_set = self._gremlin_client.submit(query, parameters)
            else:
                result_set = self._gremlin_client.submit(query)

            # Get all results
            result = result_set.all().result()

            latency_ms = (time.time() - start_time) * 1000

            # Determine if it's a write query
            query_upper = query.strip().upper()
            is_write = any(x in query_upper for x in ["ADD_V", "ADD_E", "DROP", "DELETE", "PROPERTY"])

            self.usage_stats.record_query(is_write=is_write, latency_ms=latency_ms, success=True)

            return result

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

    def execute_gremlin_batch(
        self,
        queries: List[str],
    ) -> List[List[Any]]:
        """Execute multiple Gremlin queries in batch."""
        results = []
        for query in queries:
            try:
                result = self.execute_gremlin(query)
                results.append(result)
            except Exception as e:
                self.logger.error(f"Query failed: {query}, Error: {e}")
                results.append([])
        return results

    # ---------------- SPARQL Query Execution ----------------
    def execute_sparql(
        self,
        query: str,
        use_cache: bool = False,
    ) -> List[Dict[str, Any]]:
        """Execute a SPARQL query using Neptune Data API."""
        start_time = time.time()

        try:
            client = self._get_neptune_client()

            # Execute SPARQL query
            response = client.execute_sparql_query(
                queryString=query,
                queryMode="SDB"
            )

            latency_ms = (time.time() - start_time) * 1000
            query_upper = query.strip().upper()
            is_write = any(x in query_upper for x in ["INSERT", "DELETE", "UPDATE", "CREATE", "DROP"])

            self.usage_stats.record_query(is_write=is_write, latency_ms=latency_ms, success=True)

            # Parse results
            results = []
            if "results" in response:
                for binding in response["results"].get("bindings", []):
                    result = {}
                    for key, value in binding.items():
                        result[key] = value.get("value")
                    results.append(result)

            return results

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

    # ---------------- Generic Query Execution ----------------
    def execute(
        self,
        query: str,
        parameters: Optional[Dict[str, Any]] = None,
        use_cache: bool = False,
    ) -> Any:
        """Execute a query based on the configured query mode."""
        if self.config.query_mode == "gremlin":
            return self.execute_gremlin(query, parameters, use_cache)
        elif self.config.query_mode == "sparql":
            return self.execute_sparql(query, use_cache)
        else:
            raise ValueError(f"Unsupported query mode: {self.config.query_mode}")

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

    def execute_write(
        self,
        query: str,
        parameters: Optional[Dict[str, Any]] = None,
    ) -> Any:
        """Execute a write query."""
        result = self.execute(query, parameters)
        self._cache.clear()
        return result

    # ---------------- Gremlin Vertex Operations ----------------
    def add_vertex(
        self,
        label: str,
        properties: Dict[str, Any],
    ) -> Any:
        """Add a vertex with label and properties."""
        query = f"g.addV('{label}')"
        for key, value in properties.items():
            if isinstance(value, str):
                query += f".property('{key}', '{value}')"
            else:
                query += f".property('{key}', {value})"

        result = self.execute_gremlin(query)
        self._cache.clear()
        return result[0] if result else None

    def add_vertex_eager(
        self,
        label: str,
        properties: Dict[str, Any],
    ) -> Dict[str, Any]:
        """Add a vertex and return it with all properties."""
        query = f"g.addV('{label}')"
        for key, value in properties.items():
            if isinstance(value, str):
                query += f".property('{key}', '{value}')"
            else:
                query += f".property('{key}', {value})"
        query += ".valueMap(true)"

        result = self.execute_gremlin(query)
        self._cache.clear()
        return result[0] if result else {}

    def get_vertex(
        self,
        label: str,
        properties: Dict[str, Any],
    ) -> Optional[Any]:
        """Get a vertex by label and properties."""
        conditions = " AND ".join([f"it.value('{k}') == '{v}'" for k, v in properties.items()])
        query = f"g.V().hasLabel('{label}').filter({conditions})"

        result = self.execute_gremlin(query)
        return result[0] if result else None

    def get_vertices(
        self,
        label: str,
        limit: int = 100,
    ) -> List[Any]:
        """Get all vertices with a given label."""
        query = f"g.V().hasLabel('{label}').limit({limit})"
        return self.execute_gremlin(query)

    def update_vertex(
        self,
        label: str,
        match_properties: Dict[str, Any],
        update_properties: Dict[str, Any],
    ) -> int:
        """Update vertex properties."""
        conditions = " AND ".join([f"it.value('{k}') == '{v}'" for k, v in match_properties.items()])
        query = f"g.V().hasLabel('{label}').filter({conditions})"

        for key, value in update_properties.items():
            if isinstance(value, str):
                query += f".property('{key}', '{value}')"
            else:
                query += f".property('{key}', {value})"

        result = self.execute_gremlin(query)
        self._cache.clear()
        return len(result)

    def delete_vertex(
        self,
        label: str,
        properties: Dict[str, Any],
    ) -> int:
        """Delete vertices matching the given properties."""
        conditions = " AND ".join([f"it.value('{k}') == '{v}'" for k, v in properties.items()])
        query = f"g.V().hasLabel('{label}').filter({conditions}).drop()"

        self.execute_gremlin(query)
        self._cache.clear()
        return 1

    def delete_all_vertices(self, label: str) -> int:
        """Delete all vertices with a given label."""
        query = f"g.V().hasLabel('{label}').drop()"
        self.execute_gremlin(query)
        self._cache.clear()
        return 1

    # ---------------- Gremlin Edge Operations ----------------
    def add_edge(
        self,
        source_label: str,
        source_properties: Dict[str, Any],
        target_label: str,
        target_properties: Dict[str, Any],
        edge_label: str,
        edge_properties: Optional[Dict[str, Any]] = None,
    ) -> Any:
        """Add an edge between two vertices."""
        # Find source
        src_conditions = " AND ".join([f"it.value('{k}') == '{v}'" for k, v in source_properties.items()])
        src_query = f"g.V().hasLabel('{source_label}').filter({src_conditions})"

        # Find target
        tgt_conditions = " AND ".join([f"it.value('{k}') == '{v}'" for k, v in target_properties.items()])
        tgt_query = f".V().hasLabel('{target_label}').filter({tgt_conditions})"

        # Build edge query
        query = f"{src_query}.as('a'){tgt_query}.as('b').addE('{edge_label}').from('a').to('b')"

        if edge_properties:
            for key, value in edge_properties.items():
                if isinstance(value, str):
                    query += f".property('{key}', '{value}')"
                else:
                    query += f".property('{key}', {value})"

        result = self.execute_gremlin(query)
        self._cache.clear()
        return result[0] if result else None

    def get_edges(
        self,
        edge_label: str,
        limit: int = 100,
    ) -> List[Any]:
        """Get all edges with a given label."""
        query = f"g.E().hasLabel('{edge_label}').limit({limit})"
        return self.execute_gremlin(query)

    def delete_edge(
        self,
        edge_label: str,
        properties: Dict[str, Any],
    ) -> int:
        """Delete edges matching the given properties."""
        conditions = " AND ".join([f"it.value('{k}') == '{v}'" for k, v in properties.items()])
        query = f"g.E().hasLabel('{edge_label}').filter({conditions}).drop()"

        self.execute_gremlin(query)
        self._cache.clear()
        return 1

    # ---------------- Gremlin Path Operations ----------------
    def find_shortest_path(
        self,
        source_label: str,
        source_properties: Dict[str, Any],
        target_label: str,
        target_properties: Dict[str, Any],
        max_depth: int = 5,
    ) -> List[Any]:
        """Find shortest path between two vertices."""
        src_conditions = " AND ".join([f"it.value('{k}') == '{v}'" for k, v in source_properties.items()])
        tgt_conditions = " AND ".join([f"it.value('{k}') == '{v}'" for k, v in target_properties.items()])

        query = f"""
        g.V().hasLabel('{source_label}').filter({src_conditions}).as('a')
        .V().hasLabel('{target_label}').filter({tgt_conditions}).as('b')
        .path().by('id')
        """

        return self.execute_gremlin(query)

    # ---------------- Gremlin Traversal Helpers ----------------
    def traverse(
        self,
        start_label: str,
        start_properties: Dict[str, Any],
    ) -> Any:
        """Start a traversal from a specific vertex."""
        conditions = " AND ".join([f"it.value('{k}') == '{v}'" for k, v in start_properties.items()])
        query = f"g.V().hasLabel('{start_label}').filter({conditions})"
        return TraversalHelper(self, query)

    def get_neighbors(
        self,
        vertex_id: str,
        edge_label: Optional[str] = None,
        direction: str = "both",
        limit: int = 100,
    ) -> List[Any]:
        """Get neighbors of a vertex."""
        if direction == "out":
            query = f"g.V('{vertex_id}').out()"
        elif direction == "in":
            query = f"g.V('{vertex_id}').in()"
        else:
            query = f"g.V('{vertex_id}').both()"

        if edge_label:
            query = f"g.V('{vertex_id}').{direction}('{edge_label}')"

        query += f".limit({limit})"
        return self.execute_gremlin(query)

    def get_degree(
        self,
        vertex_id: str,
        edge_label: Optional[str] = None,
    ) -> int:
        """Get the degree of a vertex."""
        if edge_label:
            query = f"g.V('{vertex_id}').bothE('{edge_label}').count()"
        else:
            query = f"g.V('{vertex_id}').bothE().count()"

        result = self.execute_gremlin(query)
        return result[0] if result else 0

    # ---------------- Schema Operations ----------------
    def get_labels(self) -> List[str]:
        """Get all vertex labels."""
        query = "g.V().label().dedup()"
        try:
            result = self.execute_gremlin(query)
            return list(set(result))
        except:
            return []

    def get_edge_labels(self) -> List[str]:
        """Get all edge labels."""
        query = "g.E().label().dedup()"
        try:
            result = self.execute_gremlin(query)
            return list(set(result))
        except:
            return []

    def get_property_keys(self) -> List[str]:
        """Get all property keys."""
        query = "g.V().properties().key().dedup()"
        try:
            result = self.execute_gremlin(query)
            return list(set(result))
        except:
            return []

    # ---------------- Utility Methods ----------------
    def count_vertices(self, label: str) -> int:
        """Count vertices with a given label."""
        query = f"g.V().hasLabel('{label}').count()"
        result = self.execute_gremlin(query)
        return result[0] if result else 0

    def count_edges(self, label: str) -> int:
        """Count edges with a given label."""
        query = f"g.E().hasLabel('{label}').count()"
        result = self.execute_gremlin(query)
        return result[0] if result else 0

    def clear_graph(self) -> bool:
        """Clear all vertices and edges from the graph."""
        query = "g.V().drop()"
        try:
            self.execute_gremlin(query)
            self._cache.clear()
            return True
        except Exception as e:
            self.logger.error(f"Clear graph 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}")


class TraversalHelper:
    """Helper class for building Gremlin traversals."""

    def __init__(self, helper: NeptuneHelper, query: str):
        self.helper = helper
        self.query = query

    def out(self, edge_label: str = None) -> "TraversalHelper":
        if edge_label:
            self.query += f".out('{edge_label}')"
        else:
            self.query += ".out()"
        return self

    def in_(self, edge_label: str = None) -> "TraversalHelper":
        if edge_label:
            self.query += f".in('{edge_label}')"
        else:
            self.query += ".in()"
        return self

    def both(self, edge_label: str = None) -> "TraversalHelper":
        if edge_label:
            self.query += f".both('{edge_label}')"
        else:
            self.query += ".both()"
        return self

    def has(self, label: str, property_key: str, value: Any) -> "TraversalHelper":
        self.query += f".has('{label}', '{property_key}', '{value}')"
        return self

    def hasLabel(self, *labels) -> "TraversalHelper":
        label_str = ", ".join([f"'{l}'" for l in labels])
        self.query += f".hasLabel({label_str})"
        return self

    def limit(self, limit: int) -> "TraversalHelper":
        self.query += f".limit({limit})"
        return self

    def count(self) -> int:
        result = self.helper.execute_gremlin(self.query + ".count()")
        return result[0] if result else 0

    def toList(self) -> List[Any]:
        return self.helper.execute_gremlin(self.query)

    def toValueMap(self) -> Dict[str, Any]:
        return self.helper.execute_gremlin(self.query + ".valueMap(true)")


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


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