"""
NebulaGraph Graph Database Helper Library
A comprehensive helper class for interacting with NebulaGraph graph database.
Provides connection management, nGQL query execution, graph operations, and vertex/edge handling.
Note: NebulaGraph uses nGQL (Nebula Graph Query Language), not Cypher.
"""

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 nebula3.gclient.net import Connection, ConnectionPool
    from nebula3.gclient.Session import Session
    from nebula3.Config import PoolConfig
    from nebula3.Exception import (
        IOErrorException,
        AuthFailedException,
        SyntaxErrorException,
    )
except ImportError:
    raise ImportError("Please install nebula3-python: pip install nebula3-python")


@dataclass
class NebulaGraphConfig:
    """Configuration class for NebulaGraph helper."""
    host: str = "127.0.0.1"
    port: int = 9669
    username: str = "root"
    password: str = "password"
    space: str = "test"
    pool_config: Optional[Dict[str, Any]] = None
    enable_caching: bool = True
    cache_ttl: int = 3600
    log_level: str = "INFO"
    # Graph space settings
    replica_factor: int = 1
    charset: str = "utf8"

    def __post_init__(self):
        # Env overrides
        self.host = os.getenv("NEBULA_HOST", self.host)
        self.port = int(os.getenv("NEBULA_PORT", str(self.port)))
        self.username = os.getenv("NEBULA_USERNAME", self.username)
        self.password = os.getenv("NEBULA_PASSWORD", self.password)
        self.space = os.getenv("NEBULA_SPACE", self.space)


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

    def __init__(self, config: Optional[NebulaGraphConfig] = None):
        self.config = config or NebulaGraphConfig()
        self._pool: Optional[ConnectionPool] = None
        self._session: Optional[Session] = 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, use_space: bool = True) -> "NebulaGraphHelper":
        """Establish connection to NebulaGraph."""
        if self._pool is None:
            # Build addresses
            addresses = [(self.config.host, self.config.port)]

            # Build pool config
            pool_config_dict = self.config.pool_config or {}
            pool_config = PoolConfig(
                max_connection_pool_size=pool_config_dict.get("max_connection_pool_size", 10),
                min_connection_pool_size=pool_config_dict.get("min_connection_pool_size", 0),
                max_connection_idle_time=pool_config_dict.get("max_connection_idle_time", 30 * 60),
                max_connection_live_time=pool_config_dict.get("max_connection_live_time", 60 * 60),
                wait_time=pool_config_dict.get("wait_time", 30),
                heartbeat_interval=pool_config_dict.get("heartbeat_interval", 30 * 60),
            )

            self._pool = ConnectionPool()
            self._pool.init(addresses, pool_config)

            # Get session
            self._session = self._pool.get_session(
                self.config.username,
                self.config.password
            )

            # Use space if specified
            if use_space:
                self.use_space(self.config.space)

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

        return self

    def connect_async(self) -> "NebulaGraphHelper":
        """Connect method alias for compatibility."""
        return self.connect()

    def close(self):
        """Close the NebulaGraph connection."""
        if self._session:
            self._session.release()
            self._session = None
        if self._pool:
            self._pool.close()
            self._pool = None
        self.logger.info("NebulaGraph 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._session is None:
                self.connect()
            result = self._session.execute("RETURN 1 AS num")
            return result.error_code == 0
        except Exception as e:
            self.logger.error(f"Connectivity check failed: {e}")
            return False

    # ---------------- Space Management ----------------
    def list_spaces(self) -> List[str]:
        """List all graph spaces."""
        result = self._session.execute("SHOW SPACES")
        if result.error_code != 0:
            self.logger.error(f"Failed to list spaces: {result.error_msg()}")
            return []
        return [row.values[0].as_string() for row in result]

    def create_space(
        self,
        space_name: str,
        partition_num: int = 10,
        replica_factor: int = 1,
        charset: str = "utf8",
        vid_type: str = "FIXED_STRING(30)",
    ) -> bool:
        """Create a new graph space."""
        cypher = f"CREATE SPACE IF NOT EXISTS {space_name} (partition_num={partition_num}, replica_factor={replica_factor}, charset={charset}, vid_type={vid_type})"
        result = self._session.execute(cypher)
        if result.error_code != 0:
            self.logger.error(f"Failed to create space: {result.error_msg()}")
            return False
        self.logger.info(f"Space '{space_name}' created")
        return True

    def use_space(self, space_name: str) -> bool:
        """Switch to a different graph space."""
        result = self._session.execute(f"USE {space_name}")
        if result.error_code != 0:
            self.logger.error(f"Failed to use space: {result.error_msg()}")
            return False
        self.config.space = space_name
        return True

    def describe_space(self, space_name: Optional[str] = None) -> Dict[str, Any]:
        """Get space information."""
        space = space_name or self.config.space
        result = self._session.execute(f"DESCRIBE SPACE {space}")
        if result.error_code != 0:
            return {}

        # Parse results
        info = {"space": space}
        return info

    # ---------------- Schema Management ----------------
    def list_tags(self) -> List[str]:
        """List all tags in current space."""
        result = self._session.execute("SHOW TAGS")
        if result.error_code != 0:
            return []
        return [row.values[0].as_string() for row in result]

    def list_edge_types(self) -> List[str]:
        """List all edge types in current space."""
        result = self._session.execute("SHOW EDGES")
        if result.error_code != 0:
            return []
        return [row.values[0].as_string() for row in result]

    def create_tag(
        self,
        tag_name: str,
        properties: Dict[str, str],
        comment: Optional[str] = None,
    ) -> bool:
        """Create a tag (node type)."""
        props_str = ", ".join([f"{k} {v}" for k, v in properties.items()])
        cypher = f"CREATE TAG IF NOT EXISTS {tag_name} ({props_str})"
        if comment:
            cypher += f" COMMENT '{comment}'"

        result = self._session.execute(cypher)
        if result.error_code != 0:
            self.logger.error(f"Failed to create tag: {result.error_msg()}")
            return False
        return True

    def create_edge_type(
        self,
        edge_name: str,
        properties: Dict[str, str],
        comment: Optional[str] = None,
    ) -> bool:
        """Create an edge type."""
        props_str = ", ".join([f"{k} {v}" for k, v in properties.items()])
        cypher = f"CREATE EDGE IF NOT EXISTS {edge_name} ({props_str})"
        if comment:
            cypher += f" COMMENT '{comment}'"

        result = self._session.execute(cypher)
        if result.error_code != 0:
            self.logger.error(f"Failed to create edge type: {result.error_msg()}")
            return False
        return True

    def describe_tag(self, tag_name: str) -> List[Dict[str, Any]]:
        """Describe a tag."""
        result = self._session.execute(f"DESCRIBE TAG {tag_name}")
        if result.error_code != 0:
            return []

        # Parse column names and rows
        columns = [col.lower() for col in result.keys()]
        return [dict(zip(columns, [self._extract_value(v) for v in row.values]))
                for row in result]

    def describe_edge_type(self, edge_name: str) -> List[Dict[str, Any]]:
        """Describe an edge type."""
        result = self._session.execute(f"DESCRIBE EDGE {edge_name}")
        if result.error_code != 0:
            return []

        columns = [col.lower() for col in result.keys()]
        return [dict(zip(columns, [self._extract_value(v) for v in row.values]))
                for row in result]

    def _extract_value(self, value) -> Any:
        """Extract value from NebulaGraph data type."""
        if value is None:
            return None
        try:
            return value.as_string()
        except:
            return str(value)

    # ---------------- 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,
        ngql: str,
        parameters: Optional[Dict[str, Any]] = None,
        use_cache: bool = False,
    ) -> List[Dict[str, Any]]:
        """Execute an nGQL query and return results as list of dicts."""
        if self._session is None:
            self.connect()

        start_time = time.time()

        try:
            # Note: NebulaGraph doesn't support parameterized queries like Neo4j
            # So we'll execute directly
            result = self._session.execute(ngql)

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

            if result.error_code != 0:
                self.logger.error(f"Query failed: {result.error_msg()}")
                self.usage_stats.record_query(latency_ms=latency_ms, success=False)
                raise Exception(f"Query failed: {result.error_msg()}")

            # Determine if it's a write query
            is_write = ngql.strip().upper().startswith((
                "INSERT", "UPDATE", "DELETE", "CREATE", "DROP", "ALTER", "USE"
            ))

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

            # Parse results
            if result.row_size() == 0:
                return []

            columns = [col for col in result.keys()]
            rows = []
            for row in result:
                row_dict = {}
                for i, col in enumerate(columns):
                    row_dict[col] = self._extract_value(row.values[i])
                rows.append(row_dict)

            return rows

        except Exception as e:
            latency_ms = (time.time() - start_time) * 1000
            self.usage_stats.record_query(latency_ms=latency_ms, success=False)
            raise

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

    def execute_and_fetch(self, ngql: str) -> List[Dict[str, Any]]:
        """Execute query and fetch all results."""
        return self._execute_query(ngql)

    # ---------------- Vertex (Node) Operations ----------------
    def insert_vertex(
        self,
        tag_name: str,
        properties: Dict[str, Any],
        vid: Optional[str] = None,
    ) -> Optional[str]:
        """Insert a vertex (node)."""
        if vid is None:
            # Generate a random VID
            import uuid
            vid = str(uuid.uuid4())

        props_str = ", ".join([f"{k}: ${k}" for k in properties.keys()])
        ngql = f"INSERT VERTEX {tag_name} ({props_str}) VALUES $vid:${props_str}"

        # Create parameter map
        params = {**properties, "vid": vid}

        result = self._execute_query(ngql, parameters=params)
        self._cache.clear()  # Clear cache after write

        return vid if result is not None else None

    def insert_vertices(
        self,
        tag_name: str,
        vertices: List[Dict[str, Any]],
    ) -> List[str]:
        """Insert multiple vertices."""
        if not vertices:
            return []

        # Build batch insert
        values_parts = []
        params = {}

        for i, v in enumerate(vertices):
            vid = v.get("vid") or f"vid_{i}"
            props = {k: v for k, v in v.items() if k != "vid"}
            props_str = ", ".join([f"{k}: ${k}_{i}" for k in props.keys()])
            values_parts.append(f"$vid_{i}:({props_str})")

            for k, val in props.items():
                params[f"{k}_{i}"] = val
            params[f"vid_{i}"] = vid

        props_sample = vertices[0]
        props_names = [k for k in props_sample.keys() if k != "vid"]
        props_str = ", ".join([f"{k} {type(val).__name__}" for k, val in props_sample.items() if k != "vid"])

        ngql = f"INSERT VERTEX {tag_name} ({props_str}) VALUES {', '.join(values_parts)}"

        result = self._execute_query(ngql, parameters=params)
        self._cache.clear()

        return [v.get("vid", f"vid_{i}") for i, v in enumerate(vertices)]

    def upsert_vertex(
        self,
        tag_name: str,
        vid: str,
        properties: Dict[str, Any],
    ) -> bool:
        """Upsert a vertex (update or insert)."""
        props_str = ", ".join([f"{k} = ${k}" for k in properties.keys()])
        ngql = f"UPSERT VERTEX {vid} SET {props_str}"

        try:
            self._execute_query(ngql, parameters=properties)
            self._cache.clear()
            return True
        except Exception as e:
            self.logger.error(f"Upsert vertex failed: {e}")
            return False

    def get_vertex(
        self,
        vid: str,
        tag_name: Optional[str] = None,
    ) -> Optional[Dict[str, Any]]:
        """Get a vertex by ID."""
        if tag_name:
            ngql = f"FETCH PROP ON {tag_name} {vid}"
        else:
            ngql = f"FETCH PROP ON * {vid}"

        result = self.execute(ngql)
        return result[0] if result else None

    def delete_vertex(self, vid: str) -> bool:
        """Delete a vertex."""
        ngql = f"DELETE VERTEX {vid}"
        try:
            self._execute_query(ngql)
            self._cache.clear()
            return True
        except Exception as e:
            self.logger.error(f"Delete vertex failed: {e}")
            return False

    def fetch_neighbors(
        self,
        vid: str,
        edge_type: Optional[str] = None,
        limit: int = 100,
    ) -> List[Dict[str, Any]]:
        """Fetch neighbors of a vertex."""
        if edge_type:
            ngql = f"GO FROM {vid} OVER {edge_type} LIMIT {limit}"
        else:
            ngql = f"GO FROM {vid} OVER * LIMIT {limit}"

        return self.execute(ngql)

    # ---------------- Edge Operations ----------------
    def insert_edge(
        self,
        edge_name: str,
        src_vid: str,
        dst_vid: str,
        properties: Optional[Dict[str, Any]] = None,
        rank: int = 0,
    ) -> bool:
        """Insert an edge between two vertices."""
        props = properties or {}

        if props:
            props_str = ", ".join([f"{k}: ${k}" for k in props.keys()])
            params = {**props}
        else:
            props_str = ""
            params = {}

        ngql = f"INSERT EDGE {edge_name} ({props_str}) VALUES {src_vid} -> {dst_vid}@{rank}:({props_str})"

        try:
            self._execute_query(ngql, parameters=params)
            self._cache.clear()
            return True
        except Exception as e:
            self.logger.error(f"Insert edge failed: {e}")
            return False

    def upsert_edge(
        self,
        edge_name: str,
        src_vid: str,
        dst_vid: str,
        properties: Dict[str, Any],
        rank: int = 0,
    ) -> bool:
        """Upsert an edge."""
        props_str = ", ".join([f"{k} = ${k}" for k in properties.keys()])
        ngql = f"UPSERT EDGE {src_vid} -> {dst_vid}@{rank} OF {edge_name} SET {props_str}"

        try:
            self._execute_query(ngql, parameters=properties)
            self._cache.clear()
            return True
        except Exception as e:
            self.logger.error(f"Upsert edge failed: {e}")
            return False

    def get_edge(
        self,
        edge_name: str,
        src_vid: str,
        dst_vid: str,
    ) -> Optional[Dict[str, Any]]:
        """Get an edge by source and destination."""
        ngql = f"FETCH PROP ON {edge_name} {src_vid} -> {dst_vid}"
        result = self.execute(ngql)
        return result[0] if result else None

    def delete_edge(
        self,
        edge_name: str,
        src_vid: str,
        dst_vid: str,
    ) -> bool:
        """Delete an edge."""
        ngql = f"DELETE EDGE {edge_name} {src_vid} -> {dst_vid}"
        try:
            self._execute_query(ngql)
            self._cache.clear()
            return True
        except Exception as e:
            self.logger.error(f"Delete edge failed: {e}")
            return False

    # ---------------- Path Operations ----------------
    def find_shortest_path(
        self,
        src_vid: str,
        dst_vid: str,
        edge_type: Optional[str] = None,
        steps: int = 5,
    ) -> List[str]:
        """Find shortest path between two vertices."""
        if edge_type:
            ngql = f"FIND SHORTEST PATH FROM {src_vid} TO {dst_vid} OVER {edge_type}"
        else:
            ngql = f"FIND SHORTEST PATH FROM {src_vid} TO {dst_vid} OVER *"

        result = self.execute(ngql)
        if result and result[0].get("path"):
            return [result[0]["path"]]
        return []

    def find_all_paths(
        self,
        src_vid: str,
        dst_vid: str,
        edge_type: Optional[str] = None,
        steps: int = 5,
    ) -> List[str]:
        """Find all paths between two vertices."""
        if edge_type:
            ngql = f"FIND ALL PATH FROM {src_vid} TO {dst_vid} OVER {edge_type} UPTO {steps} STEPS"
        else:
            ngql = f"FIND ALL PATH FROM {src_vid} TO {dst_vid} OVER * UPTO {steps} STEPS"

        result = self.execute(ngql)
        return [r.get("path", "") for r in result if r.get("path")]

    # ---------------- Query Building Helpers ----------------
    def match_vertices(
        self,
        tag_name: str,
        properties: Optional[Dict[str, Any]] = None,
        limit: int = 100,
    ) -> List[Dict[str, Any]]:
        """Match vertices by tag and properties."""
        if properties:
            props_str = " AND ".join([f"{k} == '{v}'" for k, v in properties.items()])
            ngql = f"MATCH (n:{tag_name}) WHERE {props_str} RETURN n LIMIT {limit}"
        else:
            ngql = f"MATCH (n:{tag_name}) RETURN n LIMIT {limit}"

        return self.execute(ngql)

    def match_edges(
        self,
        edge_name: str,
        src_vid: Optional[str] = None,
        dst_vid: Optional[str] = None,
        limit: int = 100,
    ) -> List[Dict[str, Any]]:
        """Match edges."""
        if src_vid and dst_vid:
            ngql = f"MATCH (a)-[r:{edge_name}]->(b) WHERE id(a) == '{src_vid} AND id(b) == '{dst_vid}' RETURN r LIMIT {limit}"
        elif src_vid:
            ngql = f"GO FROM {src_vid} OVER {edge_name} LIMIT {limit}"
        else:
            ngql = f"MATCH (a)-[r:{edge_name}]->(b) RETURN r LIMIT {limit}"

        return self.execute(ngql)

    def count_vertices(self, tag_name: str) -> int:
        """Count vertices with a given tag."""
        ngql = f"LOOKUP ON {tag_name} YIELD count(*) AS cnt"
        result = self.execute(ngql)
        return result[0].get("cnt", 0) if result else 0

    def count_edges(self, edge_name: str) -> int:
        """Count edges of a given type."""
        ngql = f"LOOKUP ON {edge_name} YIELD count(*) AS cnt"
        result = self.execute(ngql)
        return result[0].get("cnt", 0) if result else 0

    # ---------------- Index Operations ----------------
    def create_tag_index(
        self,
        index_name: str,
        tag_name: str,
        properties: List[str],
    ) -> bool:
        """Create an index on a tag."""
        props_str = ", ".join(properties)
        ngql = f"CREATE TAG INDEX IF NOT EXISTS {index_name} ON {tag_name}({props_str})"

        try:
            self.execute(ngql)
            return True
        except Exception as e:
            self.logger.error(f"Create tag index failed: {e}")
            return False

    def create_edge_index(
        self,
        index_name: str,
        edge_name: str,
        properties: List[str],
    ) -> bool:
        """Create an index on an edge type."""
        props_str = ", ".join(properties)
        ngql = f"CREATE EDGE INDEX IF NOT EXISTS {index_name} ON {edge_name}({props_str})"

        try:
            self.execute(ngql)
            return True
        except Exception as e:
            self.logger.error(f"Create edge index failed: {e}")
            return False

    def rebuild_indexes(self) -> bool:
        """Rebuild all indexes."""
        ngql = "REBUILD TAG INDEX"

        try:
            self.execute(ngql)
            return True
        except Exception as e:
            self.logger.error(f"Rebuild indexes failed: {e}")
            return False

    # ---------------- Utility Methods ----------------
    def clear_space(self, space_name: Optional[str] = None) -> bool:
        """Clear all data in a space."""
        space = space_name or self.config.space
        ngql = f"CLEAR SPACE {space}"

        try:
            self.execute(ngql)
            self._cache.clear()
            return True
        except Exception as e:
            self.logger.error(f"Clear space failed: {e}")
            return False

    def show_config(self) -> Dict[str, Any]:
        """Show current configuration."""
        return {
            "host": self.config.host,
            "port": self.config.port,
            "space": self.config.space,
            "username": self.config.username,
        }

    # ---------------- 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_tag_properties(**kwargs) -> Dict[str, str]:
    """Create tag property definitions (name -> type)."""
    return kwargs


def create_edge_properties(**kwargs) -> Dict[str, str]:
    """Create edge property definitions (name -> type)."""
    return kwargs
