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

try:
    from graphiti_core.llm_client.gemini_client import GeminiClient
    from graphiti_core.llm_client.errors import RateLimitError
    import logging
    import random

    logger = logging.getLogger(__name__)

    class HelixGeminiClient(GeminiClient):
        """
        Specialized Gemini client for Helix that adds rate limit resilience.
        
        Features:
        - Global semaphore to limit concurrency
        - Exponential backoff for RateLimitErrors (which native client ignores)
        """
        # Static semaphore shared across all instances to prevent global quota exhaustion
        _semaphore = asyncio.Semaphore(1) # Limit to 1 for maximum safety on free tier

        def __init__(self, *args, **kwargs):
            super().__init__(*args, **kwargs)
            # Ensure small_model is set from config to prevent fallback 
            # to gemini-1.5-flash or gemini-2.5-flash-lite
            if hasattr(self, 'config') and hasattr(self.config, 'small_model'):
                self.small_model = self.config.small_model
            else:
                self.small_model = getattr(self, 'model', None)

        async def generate_response(self, *args, **kwargs):
            """Override with semaphore and retry logic."""
            print(f"DEBUG: HelixGeminiClient.generate_response ENTERED (Semaphore: {self._semaphore._value})")
            async with self._semaphore:
                max_retries = 5 # 5 retries for 429
                retry_count = 0
                base_delay = 10 # Start with 10s delay
                
                while retry_count < max_retries:
                    try:
                        return await super().generate_response(*args, **kwargs)
                    except RateLimitError as e:
                        retry_count += 1
                        if retry_count >= max_retries:
                            print(f"CRITICAL: Gemini Rate Limit exhausted after {max_retries} retries.")
                            raise e
                        
                        delay = base_delay * (2 ** (retry_count - 1)) + random.uniform(0, 2)
                        print(f"WARNING: Gemini Rate Limit hit. Waiting {delay:.2f}s before retry {retry_count}/{max_retries}...")
                        await asyncio.sleep(delay)
                    except Exception as e:
                        # Capture unexpected errors in print as well
                        if "429" in str(e) or "resource_exhausted" in str(e).lower():
                             retry_count += 1
                             if retry_count >= max_retries:
                                 raise e
                             delay = base_delay * (2 ** (retry_count - 1)) + random.uniform(0, 2)
                             print(f"WARNING: Gemini 429 (detected as generic exception). Waiting {delay:.2f}s...")
                             await asyncio.sleep(delay)
                             continue
                        raise e

except ImportError:
    # If graphiti-core is not installed yet
    class HelixGeminiClient:
        pass

@dataclass
class GraphitiStorage(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")
    )
    
    # Clients for Graphiti (can be instances or None)
    llm_client: Any = None
    embedder: Any = None
    cross_encoder: Any = None

    llm_model_name: str | None = None
    llm_api_key: str | None = None
    embedding_dim: int = 1536
    
    _initialized: bool = False
    
    async def initialize(self) -> None:
        """Initialize Graphiti connection and build indices."""
        if self._initialized:
            return
            
        print(f"DEBUG: Initializing GraphitiStorage. LLM Client: {self.llm_client}, Model: {self.llm_model_name}, Dim: {self.embedding_dim}")
            
        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:
            # Lazy initialize Gemini clients if credentials provided for native support
            if self.llm_api_key and self.llm_model_name:
                if "gemini" in self.llm_model_name.lower():
                    from google import genai
                    from graphiti_core.llm_client.gemini_client import GeminiClient
                    from graphiti_core.llm_client.config import LLMConfig
                    from graphiti_core.embedder.gemini import GeminiEmbedder, GeminiEmbedderConfig
                    from graphiti_core.cross_encoder.gemini_reranker_client import GeminiRerankerClient
                    
                    # Create shared genai.Client
                    # Must use v1beta — v1 doesn't support systemInstruction,
                    # responseMimeType, or responseSchema fields that Graphiti's
                    # GeminiClient._generate_response() relies on.
                    genai_client = genai.Client(api_key=self.llm_api_key, http_options={'api_version': 'v1beta'})
                    
                    # Configure native Graphiti clients
                    llm_config = LLMConfig(
                        api_key=self.llm_api_key, 
                        model=self.llm_model_name,
                        small_model=self.llm_model_name
                    )
                    embed_config = GeminiEmbedderConfig(
                        embedding_model=os.getenv("EMBEDDING_MODEL_NAME", "gemini-embedding-001"), 
                        api_key=self.llm_api_key,
                        embedding_dim=self.embedding_dim
                    )
                    
                    if self.llm_client is None:
                        print(f"DEBUG: Creating HelixGeminiClient with model {self.llm_model_name}")
                        self.llm_client = HelixGeminiClient(config=llm_config, client=genai_client)
                    else:
                        print(f"DEBUG: Using existing llm_client: {self.llm_client}")
                    if self.embedder is None:
                        self.embedder = GeminiEmbedder(config=embed_config, client=genai_client)
                    if self.cross_encoder is None:
                        # GeminiRerankerClient uses LLMConfig
                        self.cross_encoder = GeminiRerankerClient(config=llm_config, client=genai_client)
            
            # Use provided clients or default to OpenAI (legacy/bridge mode)
            self.graphiti = Graphiti(
                uri=self.neo4j_uri,
                user=self.neo4j_user,
                password=self.neo4j_password,
                llm_client=self.llm_client,
                embedder=self.embedder,
                cross_encoder=self.cross_encoder
            )
            print(f"DEBUG: Graphiti initialized with LLM Client: {self.graphiti.llm_client}")
        
        # Build indices and constraints for optimal performance
        # Force recreation if specifically requested via env (useful for dimension changes)
        delete_existing = os.getenv("HELIX_REBUILD_INDICES", "false").lower() == "true"
        if delete_existing:
             print("DEBUG: Rebuilding Graphiti indices and constraints...")
        await self.graphiti.build_indices_and_constraints(delete_existing=delete_existing)
        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:
        """
        Called after LightRAG finishes indexing a batch of documents.
        
        All node and edge writes are done immediately in upsert_node()
        and upsert_edge(), so this is intentionally a no-op.
        
        We do NOT call graphiti.add_episode() here: it triggers ~10
        extra LLM calls per entity (entity resolution, edge extraction,
        summary generation, etc.) and Graphiti's entity resolution can
        rename Entity nodes, breaking the name consistency between
        LightRAG's vector DB and Neo4j.
        """
        pass
    
    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"}
    
    # ============================================================
    # Neo4j Result Helpers
    # ============================================================
    
    @staticmethod
    def _records(result) -> list:
        """Extract records list from execute_query result.
        
        neo4j's execute_query returns EagerResult(records, summary, keys).
        EagerResult is a NamedTuple that is ALWAYS truthy, even with
        zero records.  result[0] gives the records *list*, not a single
        Record.  This helper normalises both EagerResult and plain list
        return values so callers can simply check ``if records:``.
        """
        if result is None:
            return []
        # EagerResult has a .records attribute
        if hasattr(result, 'records'):
            return result.records or []
        # Fallback: if first element is itself a list, treat it as
        # EagerResult-like (records, summary, keys)
        try:
            first = result[0]
            if isinstance(first, list):
                return first
        except (IndexError, TypeError):
            pass
        # Assume result is already a records list
        return list(result) if result else []
    
    @staticmethod
    def _to_dict(record) -> dict:
        """Convert a neo4j Record to a plain dict."""
        if hasattr(record, 'data'):
            return record.data()
        if hasattr(record, 'keys'):
            return {k: record[k] for k in record.keys()}
        return dict(record)
    
    # ============================================================
    # 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,
            )
            records = self._records(result)
            return records[0]["exists"] if records else False
        except Exception:
            return False
    
    async def get_node(self, node_id: str) -> dict[str, str] | None:
        """Get node properties by ID.
        
        Maps between Graphiti's Neo4j schema and LightRAG's expected format:
        - Graphiti stores: name, summary, labels
        - LightRAG expects: entity_id, entity_type, description, source_id
        
        We read both the LightRAG metadata we wrote (source_id, description,
        entity_type) and Graphiti's native properties (summary) as fallback.
        """
        if not self._initialized:
            await self.initialize()
            
        try:
            query = """
            MATCH (n:Entity {name: $name})
            RETURN n.name as entity_id,
                   COALESCE(n.entity_type, '') as entity_type,
                   COALESCE(n.description, n.summary, '') as description,
                   COALESCE(n.source_id, '') as source_id,
                   COALESCE(n.file_path, '') as file_path
            LIMIT 1
            """
            result = await self.graphiti.driver.execute_query(
                query,
                name=node_id,
            )
            records = self._records(result)
            if records:
                row = self._to_dict(records[0])
                # Ensure entity_name is set (LightRAG uses this in context assembly)
                row["entity_name"] = row.get("entity_id", node_id)
                return row
            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 directly in Neo4j.
        
        Uses MERGE to create the Entity node if it doesn't exist,
        then writes all LightRAG-required properties (entity_type,
        description, source_id, file_path) directly onto the node.
        
        This bypasses Graphiti's add_episode() to avoid:
        - ~10 extra LLM calls per entity (entity resolution, etc.)
        - Entity renaming that breaks LightRAG vector DB consistency
        """
        if not self._initialized:
            await self.initialize()
        
        source_id = node_data.get("source_id", "")
        description = node_data.get("description", "")
        entity_type = node_data.get("entity_type", "")
        file_path = node_data.get("file_path", "")
        created_at = node_data.get("created_at", int(datetime.now(timezone.utc).timestamp()))
        
        try:
            merge_query = """
            MERGE (n:Entity {name: $name})
            ON CREATE SET
                n.entity_type = $entity_type,
                n.description = $description,
                n.source_id   = $source_id,
                n.file_path   = $file_path,
                n.created_at  = datetime(),
                n.group_id    = $group_id
            ON MATCH SET
                n.entity_type = CASE
                    WHEN n.entity_type IS NULL OR n.entity_type = '' THEN $entity_type
                    ELSE n.entity_type
                  END,
                n.description = CASE
                    WHEN n.description IS NULL OR n.description = '' THEN $description
                    WHEN $description = '' THEN n.description
                    ELSE n.description + ' ' + $description
                  END,
                n.source_id = CASE
                    WHEN n.source_id IS NULL OR n.source_id = '' THEN $source_id
                    WHEN $source_id = '' THEN n.source_id
                    ELSE n.source_id + $sep + $source_id
                  END,
                n.file_path = CASE
                    WHEN n.file_path IS NULL OR n.file_path = '' THEN $file_path
                    ELSE n.file_path
                  END
            """
            await self.graphiti.driver.execute_query(
                merge_query,
                name=node_id,
                entity_type=entity_type,
                description=description,
                source_id=source_id,
                file_path=file_path,
                group_id=self.namespace or "default",
                sep="<SEP>",
            )
        except Exception as e:
            print(f"Warning: Failed to upsert node {node_id}: {e}")
    
    async def delete_node(self, node_id: str) -> None:
        """Delete a node from the graph."""
        if not self._initialized:
            await self.initialize()
        try:
            await self.graphiti.driver.execute_query(
                "MATCH (n:Entity {name: $name}) DETACH DELETE n",
                name=node_id,
            )
        except Exception:
            pass
        
    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,
            )
            records = self._records(result)
            return records[0]["exists"] if records 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.
        
        Maps Graphiti's RELATES_TO edge properties to LightRAG's expected format:
        - Graphiti stores: fact, name, valid_at, invalid_at, episodes
        - LightRAG expects: description, weight, keywords, source_id
        """
        if not self._initialized:
            await self.initialize()
            
        try:
            query = """
            MATCH (a:Entity {name: $source})-[r]-(b:Entity {name: $target})
            RETURN COALESCE(r.description, r.fact, r.name, '') as description,
                   COALESCE(r.weight, 1.0) as weight,
                   COALESCE(r.keywords, type(r), '') as keywords,
                   COALESCE(r.source_id, '') as source_id
            LIMIT 1
            """
            result = await self.graphiti.driver.execute_query(
                query,
                source=source_node_id,
                target=target_node_id,
            )
            records = self._records(result)
            if records:
                row = self._to_dict(records[0])
                # Ensure weight is numeric
                if row.get("weight") is None:
                    row["weight"] = 1.0
                return row
            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 directly in Neo4j.
        
        Ensures both endpoint Entity nodes exist first (MERGE),
        then creates a RELATES_TO relationship with LightRAG-
        compatible properties.
        """
        if not self._initialized:
            await self.initialize()
        
        description = edge_data.get("description", "")
        keywords = edge_data.get("keywords", "")
        weight = edge_data.get("weight", 1.0)
        source_id = edge_data.get("source_id", "")
        
        try:
            merge_query = """
            MERGE (a:Entity {name: $source})
            MERGE (b:Entity {name: $target})
            MERGE (a)-[r:RELATES_TO]->(b)
            ON CREATE SET
                r.description = $description,
                r.keywords    = $keywords,
                r.weight      = $weight,
                r.source_id   = $source_id,
                r.created_at  = datetime()
            ON MATCH SET
                r.description = CASE
                    WHEN r.description IS NULL OR r.description = '' THEN $description
                    WHEN $description = '' THEN r.description
                    ELSE r.description + ' ' + $description
                  END,
                r.weight   = $weight,
                r.source_id = CASE
                    WHEN r.source_id IS NULL OR r.source_id = '' THEN $source_id
                    WHEN $source_id = '' THEN r.source_id
                    ELSE r.source_id + $sep + $source_id
                  END
            """
            await self.graphiti.driver.execute_query(
                merge_query,
                source=source_node_id,
                target=target_node_id,
                description=description,
                keywords=keywords,
                weight=float(weight) if weight else 1.0,
                source_id=source_id,
                sep="<SEP>",
            )
        except Exception as e:
            print(f"Warning: Failed to upsert edge {source_node_id}->{target_node_id}: {e}")
    
    async def remove_edges(self, edges: list[tuple[str, str]]) -> None:
        """Delete multiple edges."""
        if not self._initialized:
            await self.initialize()
        for src, tgt in edges:
            try:
                await self.graphiti.driver.execute_query(
                    "MATCH (a:Entity {name: $source})-[r:RELATES_TO]-(b:Entity {name: $target}) DELETE r",
                    source=src,
                    target=tgt,
                )
            except Exception:
                pass
    
    # ============================================================
    # 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,
            )
            records = self._records(result)
            return records[0]["degree"] if records 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,
            )
            records = self._records(results)
            if records:
                return [(r["source"], r["target"]) for r in records]
            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)
            records = self._records(results)
            return [r["name"] for r in records] if records 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,
            )
            records = self._records(results)
            return [r["name"] for r in records] if records 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,
            )
            records = self._records(results)
            return [r["name"] for r in records] if records 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,
                   COALESCE(n.entity_type, '') as entity_type,
                   COALESCE(n.description, n.summary, '') as description,
                   COALESCE(n.source_id, '') as source_id
            """
            results = await self.graphiti.driver.execute_query(query)
            records = self._records(results)
            return [self._to_dict(r) for r in records] if records 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)
            records = self._records(results)
            return [self._to_dict(r) for r in records] if records 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)
