"""
GraphitiGraphStorage: Custom graph storage adapter for LightRAG using Graphiti.

This module implements LightRAG's BaseGraphStorage interface using Graphiti
as the backend, enabling bi-temporal Knowledge Graph capabilities.
"""

from __future__ import annotations

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

from dotenv import load_dotenv

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

# Import base class from lightrag
try:
    from lightrag.base import BaseGraphStorage
    from lightrag.types import KnowledgeGraph
    from lightrag.utils import EmbeddingFunc
except ImportError:
    # Fallback for development without lightrag installed
    from dataclasses import dataclass as base_dataclass
    
    @base_dataclass
    class BaseGraphStorage:
        namespace: str = ""
        workspace: str = ""
        global_config: dict = field(default_factory=dict)
        embedding_func: Any = None
        
    class KnowledgeGraph:
        pass
        
    class EmbeddingFunc:
        pass


@dataclass
class GraphitiGraphStorage(BaseGraphStorage):
    """
    Graph storage adapter that bridges LightRAG to Graphiti's temporal KG.
    
    This implementation:
    - Uses Graphiti's add_episode() for entity/relation storage
    - Tracks bi-temporal metadata (t_valid, t_created)
    - Leverages Graphiti's automatic entity resolution
    - Supports incremental updates without full graph rebuild
    
    Configuration via environment variables:
    - NEO4J_URI: Neo4j connection URI
    - NEO4J_USERNAME: Neo4j username
    - NEO4J_PASSWORD: Neo4j password
    - NEO4J_DATABASE: Neo4j database name
    """
    
    # Graphiti instance (injected or created)
    graphiti: Any = None
    
    # Neo4j configuration
    neo4j_uri: str = field(
        default_factory=lambda: os.getenv("NEO4J_URI", "bolt://localhost:7687")
    )
    neo4j_user: str = field(
        default_factory=lambda: os.getenv("NEO4J_USERNAME", "neo4j")
    )
    neo4j_password: str = field(
        default_factory=lambda: os.getenv("NEO4J_PASSWORD", "password")
    )
    neo4j_database: str = field(
        default_factory=lambda: os.getenv("NEO4J_DATABASE", "neo4j")
    )
    
    # Cache for pending operations
    _pending_nodes: dict = field(default_factory=dict)
    _pending_edges: list = field(default_factory=list)
    _initialized: bool = False
    
    async def initialize(self) -> None:
        """Initialize Graphiti connection and build indices."""
        if self._initialized:
            return
            
        try:
            from graphiti_core import Graphiti
        except ImportError:
            raise ImportError(
                "graphiti-core is required. Install with: pip install graphiti-core"
            )
        
        if self.graphiti is None:
            self.graphiti = Graphiti(
                uri=self.neo4j_uri,
                user=self.neo4j_user,
                password=self.neo4j_password,
            )
        
        # Build indices and constraints for optimal performance
        await self.graphiti.build_indices_and_constraints()
        self._initialized = True
        
    async def finalize(self) -> None:
        """Close Graphiti connection and cleanup."""
        if self.graphiti:
            await self.graphiti.close()
            self.graphiti = None
        self._initialized = False
        
    async def index_done_callback(self) -> None:
        """
        Commit pending operations after indexing batch.
        
        This is called by LightRAG after a batch of upsert operations.
        We flush any pending episodes to Graphiti here.
        """
        if not self._pending_nodes and not self._pending_edges:
            return
            
        # Process pending nodes as episodes
        from graphiti_core.nodes import EpisodeType
        
        for node_id, node_data in self._pending_nodes.items():
            try:
                await self.graphiti.add_episode(
                    name=f"entity_{node_id}",
                    episode_body=self._format_entity_episode(node_id, node_data),
                    source=EpisodeType.text,
                    source_description=node_data.get("source_id", "lightrag"),
                    reference_time=self._get_reference_time(node_data),
                    group_id=self.namespace or "default",
                )
            except Exception as e:
                # Log but don't fail the entire batch
                print(f"Warning: Failed to add entity {node_id}: {e}")
        
        # Process pending edges as episodes
        for edge_data in self._pending_edges:
            try:
                await self.graphiti.add_episode(
                    name=f"relation_{edge_data['src_id']}_{edge_data['tgt_id']}",
                    episode_body=self._format_relation_episode(edge_data),
                    source=EpisodeType.text,
                    source_description=edge_data.get("source_id", "lightrag"),
                    reference_time=self._get_reference_time(edge_data),
                    group_id=self.namespace or "default",
                )
            except Exception as e:
                print(f"Warning: Failed to add relation: {e}")
        
        # Clear pending operations
        self._pending_nodes.clear()
        self._pending_edges.clear()
    
    def _format_entity_episode(self, entity_name: str, data: dict) -> str:
        """Format entity data as episode text for Graphiti."""
        entity_type = data.get("entity_type", "entity")
        description = data.get("description", "")
        return f"{entity_name} is a {entity_type}. {description}"
    
    def _format_relation_episode(self, data: dict) -> str:
        """Format relation data as episode text for Graphiti."""
        src = data.get("src_id", "")
        tgt = data.get("tgt_id", "")
        keywords = data.get("keywords", "relates to")
        description = data.get("description", "")
        return f"{src} {keywords} {tgt}. {description}"
    
    def _get_reference_time(self, data: dict) -> datetime:
        """Extract reference time from data or use current time."""
        if "reference_time" in data and data["reference_time"]:
            return data["reference_time"]
        if "timestamp" in data:
            try:
                return datetime.fromtimestamp(data["timestamp"], tz=timezone.utc)
            except (ValueError, TypeError):
                pass
        return datetime.now(timezone.utc)
    
    async def drop(self) -> dict[str, str]:
        """
        Drop all data from storage.
        
        Note: This is a destructive operation that clears the entire graph.
        """
        # TODO: Implement graph clearing via Graphiti/Neo4j
        return {"status": "success", "message": "data dropped"}
    
    # ============================================================
    # Node Operations
    # ============================================================
    
    async def has_node(self, node_id: str) -> bool:
        """Check if node exists in the Knowledge Graph."""
        if not self._initialized:
            await self.initialize()
            
        # Query Neo4j via Graphiti's driver
        try:
            query = """
            MATCH (n:Entity {name: $name})
            RETURN count(n) > 0 as exists
            """
            result = await self.graphiti.driver.execute_query(
                query,
                name=node_id,
            )
            return result[0]["exists"] if result else False
        except Exception:
            return False
    
    async def get_node(self, node_id: str) -> dict[str, str] | None:
        """Get node properties by ID."""
        if not self._initialized:
            await self.initialize()
            
        try:
            query = """
            MATCH (n:Entity {name: $name})
            RETURN n.name as entity_id, 
                   n.entity_type as entity_type,
                   n.description as description,
                   n.source_id as source_id
            LIMIT 1
            """
            result = await self.graphiti.driver.execute_query(
                query,
                name=node_id,
            )
            if result:
                return dict(result[0])
            return None
        except Exception:
            return None
    
    async def upsert_node(self, node_id: str, node_data: dict[str, str]) -> None:
        """
        Insert or update a node in the Knowledge Graph.
        
        Instead of directly writing to Neo4j, we queue the node data
        and write it as an episode during index_done_callback().
        This allows Graphiti to:
        - Extract additional entities from the description
        - Perform entity resolution (deduplication)
        - Track temporal metadata
        """
        self._pending_nodes[node_id] = {
            **node_data,
            "entity_id": node_id,
        }
    
    async def delete_node(self, node_id: str) -> None:
        """Delete a node from the graph."""
        # Remove from pending if exists
        if node_id in self._pending_nodes:
            del self._pending_nodes[node_id]
            
        # Note: Graphiti uses soft-delete via edge invalidation
        # Full deletion would require direct Neo4j query
        
    async def remove_nodes(self, nodes: list[str]) -> None:
        """Delete multiple nodes."""
        for node_id in nodes:
            await self.delete_node(node_id)
    
    # ============================================================
    # Edge Operations
    # ============================================================
    
    async def has_edge(self, source_node_id: str, target_node_id: str) -> bool:
        """Check if edge exists between two nodes."""
        if not self._initialized:
            await self.initialize()
            
        try:
            query = """
            MATCH (a:Entity {name: $source})-[r]-(b:Entity {name: $target})
            RETURN count(r) > 0 as exists
            """
            result = await self.graphiti.driver.execute_query(
                query,
                source=source_node_id,
                target=target_node_id,
            )
            return result[0]["exists"] if result else False
        except Exception:
            return False
    
    async def get_edge(
        self, source_node_id: str, target_node_id: str
    ) -> dict[str, str] | None:
        """Get edge properties between two nodes."""
        if not self._initialized:
            await self.initialize()
            
        try:
            query = """
            MATCH (a:Entity {name: $source})-[r]-(b:Entity {name: $target})
            RETURN r.fact as description,
                   r.weight as weight,
                   type(r) as keywords
            LIMIT 1
            """
            result = await self.graphiti.driver.execute_query(
                query,
                source=source_node_id,
                target=target_node_id,
            )
            if result:
                return dict(result[0])
            return None
        except Exception:
            return None
    
    async def upsert_edge(
        self,
        source_node_id: str,
        target_node_id: str,
        edge_data: dict[str, str],
    ) -> None:
        """
        Insert or update an edge in the Knowledge Graph.
        
        Queued for batch processing via Graphiti episodes.
        """
        self._pending_edges.append({
            "src_id": source_node_id,
            "tgt_id": target_node_id,
            **edge_data,
        })
    
    async def remove_edges(self, edges: list[tuple[str, str]]) -> None:
        """Delete multiple edges."""
        # Remove from pending
        self._pending_edges = [
            e for e in self._pending_edges
            if (e["src_id"], e["tgt_id"]) not in edges
        ]
    
    # ============================================================
    # Degree Operations
    # ============================================================
    
    async def node_degree(self, node_id: str) -> int:
        """Get the degree (number of connected edges) of a node."""
        if not self._initialized:
            await self.initialize()
            
        try:
            query = """
            MATCH (n:Entity {name: $name})-[r]-()
            RETURN count(r) as degree
            """
            result = await self.graphiti.driver.execute_query(
                query,
                name=node_id,
            )
            return result[0]["degree"] if result else 0
        except Exception:
            return 0
    
    async def edge_degree(self, src_id: str, tgt_id: str) -> int:
        """Get total degree of an edge (sum of source and target degrees)."""
        src_degree = await self.node_degree(src_id)
        tgt_degree = await self.node_degree(tgt_id)
        return src_degree + tgt_degree
    
    # ============================================================
    # Edge Retrieval Operations
    # ============================================================
    
    async def get_node_edges(
        self, source_node_id: str
    ) -> list[tuple[str, str]] | None:
        """Get all edges connected to a node."""
        if not self._initialized:
            await self.initialize()
            
        try:
            query = """
            MATCH (a:Entity {name: $name})-[r]-(b:Entity)
            RETURN a.name as source, b.name as target
            """
            results = await self.graphiti.driver.execute_query(
                query,
                name=source_node_id,
            )
            if results:
                return [(r["source"], r["target"]) for r in results]
            return None
        except Exception:
            return None
    
    # ============================================================
    # Label/Entity Operations
    # ============================================================
    
    async def get_all_labels(self) -> list[str]:
        """Get all entity names in the graph."""
        if not self._initialized:
            await self.initialize()
            
        try:
            query = """
            MATCH (n:Entity)
            RETURN DISTINCT n.name as name
            ORDER BY name
            """
            results = await self.graphiti.driver.execute_query(query)
            return [r["name"] for r in results] if results else []
        except Exception:
            return []
    
    async def get_popular_labels(self, limit: int = 300) -> list[str]:
        """Get popular entities by degree."""
        if not self._initialized:
            await self.initialize()
            
        try:
            query = """
            MATCH (n:Entity)-[r]-()
            WITH n.name as name, count(r) as degree
            ORDER BY degree DESC
            LIMIT $limit
            RETURN name
            """
            results = await self.graphiti.driver.execute_query(
                query,
                limit=limit,
            )
            return [r["name"] for r in results] if results else []
        except Exception:
            return []
    
    async def search_labels(self, query: str, limit: int = 50) -> list[str]:
        """Search entity names with fuzzy matching."""
        if not self._initialized:
            await self.initialize()
            
        try:
            cypher = """
            MATCH (n:Entity)
            WHERE toLower(n.name) CONTAINS toLower($query)
            RETURN n.name as name
            LIMIT $limit
            """
            results = await self.graphiti.driver.execute_query(
                cypher,
                query=query,
                limit=limit,
            )
            return [r["name"] for r in results] if results else []
        except Exception:
            return []
    
    # ============================================================
    # Knowledge Graph Retrieval
    # ============================================================
    
    async def get_all_nodes(self) -> list[dict]:
        """Get all nodes in the graph."""
        if not self._initialized:
            await self.initialize()
            
        try:
            query = """
            MATCH (n:Entity)
            RETURN n.name as entity_id,
                   n.entity_type as entity_type,
                   n.description as description
            """
            results = await self.graphiti.driver.execute_query(query)
            return [dict(r) for r in results] if results else []
        except Exception:
            return []
    
    async def get_all_edges(self) -> list[dict]:
        """Get all edges in the graph."""
        if not self._initialized:
            await self.initialize()
            
        try:
            query = """
            MATCH (a:Entity)-[r]->(b:Entity)
            RETURN a.name as source,
                   b.name as target,
                   type(r) as relation_type,
                   r.fact as description
            """
            results = await self.graphiti.driver.execute_query(query)
            return [dict(r) for r in results] if results else []
        except Exception:
            return []
    
    async def get_knowledge_graph(
        self,
        node_label: str,
        max_depth: int = 3,
        max_nodes: int = 1000,
    ) -> KnowledgeGraph:
        """
        Retrieve a connected subgraph starting from a node.
        
        Args:
            node_label: Starting entity name, or "*" for all.
            max_depth: Maximum traversal depth.
            max_nodes: Maximum nodes to return.
            
        Returns:
            KnowledgeGraph with nodes and edges.
        """
        if not self._initialized:
            await self.initialize()
            
        # TODO: Implement proper subgraph traversal
        # For now, return empty KnowledgeGraph
        from lightrag.types import KnowledgeGraph as KG
        
        return KG(nodes=[], edges=[], is_truncated=False)
