"""
Helix: Main orchestration class combining LightRAG with Graphiti.

This module provides the Helix class which wraps LightRAG's proven RAG 
retrieval capabilities with Graphiti's temporal Knowledge Graph production.
"""

from __future__ import annotations

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

from dotenv import load_dotenv

# Load environment variables
load_dotenv(dotenv_path=".env", override=False)


@dataclass
class HelixConfig:
    """Configuration for Helix instance."""
    
    # Working directory
    working_dir: str = field(default="./helix_storage")
    
    # Neo4j configuration (for Graphiti)
    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")
    )
    
    # Supabase configuration (for vector storage)
    supabase_url: str = field(
        default_factory=lambda: os.getenv("SUPABASE_URL", "")
    )
    supabase_key: str = field(
        default_factory=lambda: os.getenv("SUPABASE_KEY", "")
    )
    
    # LLM configuration (generic - no provider specified)
    llm_model_name: str = field(
        default_factory=lambda: os.getenv("LLM_MODEL_NAME", "")
    )
    llm_api_key: str = field(
        default_factory=lambda: os.getenv("LLM_API_KEY", "")
    )
    llm_api_base: str = field(
        default_factory=lambda: os.getenv("LLM_API_BASE", "")
    )
    
    # Embedding configuration
    embedding_model_name: str = field(
        default_factory=lambda: os.getenv("EMBEDDING_MODEL_NAME", "")
    )
    embedding_dim: int = field(
        default_factory=lambda: int(os.getenv("EMBEDDING_DIM", "1536"))
    )


class Helix:
    """
    Helix: LightRAG + Graphiti for Temporal GraphRAG.
    
    Combines LightRAG's proven dual-level retrieval (low + high) with
    Graphiti's bi-temporal Knowledge Graph for:
    - Temporal query support (point-in-time queries)
    - Automatic edge invalidation (conflict resolution)
    - Incremental updates (no full graph rebuild)
    - Graph-aligned hallucination detection
    
    Example:
        ```python
        from helix import Helix
        
        async def main():
            helix = Helix()
            await helix.initialize()
            
            # Insert document with temporal tracking
            await helix.insert(
                "Alan Turing was born on June 23, 1912.",
                reference_time=datetime(1912, 6, 23)
            )
            
            # Query with temporal awareness
            result = await helix.query(
                "When was Alan Turing born?",
                mode="hybrid"
            )
            print(result)
            
            await helix.close()
        ```
    """
    
    def __init__(
        self,
        config: HelixConfig | None = None,
        working_dir: str | None = None,
        llm_model_func: Callable | None = None,
        embedding_func: Callable | None = None,
    ):
        """
        Initialize Helix instance.
        
        Args:
            config: HelixConfig instance with all settings. If None, uses defaults from env.
            working_dir: Override working directory (takes precedence over config).
            llm_model_func: Custom LLM function for inference.
            embedding_func: Custom embedding function.
        """
        self.config = config or HelixConfig()
        
        if working_dir:
            self.config.working_dir = working_dir
            
        self.llm_model_func = llm_model_func
        self.embedding_func = embedding_func
        
        # Will be initialized in initialize()
        self.graphiti = None
        self.lightrag = None
        self._initialized = False
        
    async def initialize(self) -> None:
        """
        Initialize Graphiti and LightRAG instances.
        
        Must be called before any insert/query operations.
        """
        if self._initialized:
            return
            
        # Import here to avoid circular imports and allow lazy loading
        try:
            from graphiti_core import Graphiti
            from graphiti_core.nodes import EpisodeType
        except ImportError:
            raise ImportError(
                "graphiti-core is required for Helix. "
                "Install with: pip install graphiti-core"
            )
            
        try:
            from lightrag import LightRAG
        except ImportError:
            raise ImportError(
                "lightrag is required for Helix. "
                "Install with: pip install lightrag-hku"
            )
        
        # 0. Configure Graphiti LLM/Embedder
        # Graphiti needs an LLM client for its own operations (entity resolution,
        # etc.), but since we've bypassed add_episode() entirely (index_done_callback
        # is a no-op and upsert_node/edge write directly to Neo4j), Graphiti's LLM
        # is never actually called.  For non-Gemini providers we can use the OpenAI
        # client with a dummy key — it won't be invoked at runtime.
        graphiti_kwargs = {
            "uri": self.config.neo4j_uri,
            "user": self.config.neo4j_user,
            "password": self.config.neo4j_password,
        }
        
        print(f"DEBUG: Helix.initialize - Model: {self.config.llm_model_name}, Dim: {self.config.embedding_dim}")
        
        if self.config.llm_model_name and "gemini" in self.config.llm_model_name.lower():
            # Hotfix: Set environment variables that Graphiti might pick up automatically
            os.environ["OPENAI_API_KEY"] = self.config.llm_api_key
            os.environ["OPENAI_API_BASE"] = "https://generativelanguage.googleapis.com/v1beta"
            
            try:
                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
                
                llm_config = LLMConfig(
                    api_key=self.config.llm_api_key,
                    model=self.config.llm_model_name,
                    small_model=self.config.llm_model_name
                )
                graphiti_kwargs["llm_client"] = GeminiClient(config=llm_config)
                
                embed_config = GeminiEmbedderConfig(
                    api_key=self.config.llm_api_key,
                    embedding_model=self.config.embedding_model_name,
                    embedding_dim=self.config.embedding_dim
                )
                graphiti_kwargs["embedder"] = GeminiEmbedder(config=embed_config)
                graphiti_kwargs["cross_encoder"] = GeminiRerankerClient(config=llm_config)
                
            except ImportError:
                print("[WARNING] native Gemini support for Graphiti failed to load. Falling back to OpenAI bridge.")
                from graphiti_core.llm_client import OpenAIClient, LLMConfig
                from graphiti_core.embedder import OpenAIEmbedder, OpenAIEmbedderConfig
                
                gemini_base_url = "https://generativelanguage.googleapis.com"
                llm_config = LLMConfig(
                    api_key=self.config.llm_api_key,
                    model=self.config.llm_model_name,
                    base_url=gemini_base_url
                )
                graphiti_kwargs["llm_client"] = OpenAIClient(config=llm_config)
                embed_config = OpenAIEmbedderConfig(
                    api_key=self.config.llm_api_key,
                    base_url=gemini_base_url,
                    model="gemini-embedding-001"
                )
                graphiti_kwargs["embedder"] = OpenAIEmbedder(config=embed_config)
        else:
            # Non-Gemini provider (HuggingFace, Ollama, OpenAI, etc.)
            # Graphiti requires an LLM client at init time but since we've
            # bypassed add_episode(), it's never called.  Use OpenAI client
            # with a placeholder key so Graphiti's constructor is happy.
            try:
                from graphiti_core.llm_client import OpenAIClient
                from graphiti_core.llm_client.config import LLMConfig
                
                api_key = self.config.llm_api_key or os.getenv("OPENAI_API_KEY", "sk-placeholder-not-used")
                # Also set env var so Graphiti's internal checks pass
                if not os.getenv("OPENAI_API_KEY"):
                    os.environ["OPENAI_API_KEY"] = api_key
                
                llm_config = LLMConfig(
                    api_key=api_key,
                    model=self.config.llm_model_name or "placeholder",
                    base_url=self.config.llm_api_base or None,
                )
                graphiti_kwargs["llm_client"] = OpenAIClient(config=llm_config)
            except ImportError:
                print("[WARNING] Could not configure Graphiti LLM client.")

        # Initialize Graphiti for temporal Knowledge Graph
        try:
            self.graphiti = Graphiti(**graphiti_kwargs)
        except Exception as e:
            raise RuntimeError(f"Failed to initialize Graphiti with kwargs {graphiti_kwargs.keys()}: {e}")
        
        # Build indices and constraints (with timeout to avoid hangs on Aura)
        try:
            await asyncio.wait_for(
                self.graphiti.build_indices_and_constraints(),
                timeout=30.0
            )
        except asyncio.TimeoutError:
            print("[WARNING] Graphiti index building timed out. Proceeding anyway...")
        except Exception as e:
            print(f"[WARNING] Graphiti index building failed: {e}. Proceeding anyway...")
        
        # Initialize LightRAG with custom storage
        lightrag_kwargs = {
            "working_dir": self.config.working_dir,
            "graph_storage": "GraphitiStorage",
            "llm_model_name": self.config.llm_model_name,
        }
        
        # Set Neo4j environment variables for GraphitiStorage to pick up
        os.environ["NEO4J_URI"] = self.config.neo4j_uri
        os.environ["NEO4J_USERNAME"] = self.config.neo4j_user
        os.environ["NEO4J_PASSWORD"] = self.config.neo4j_password
        
        # 1. Resolve LLM Model Function
        actual_llm_func = self.llm_model_func
        if not actual_llm_func and self.config.llm_model_name:
            mn = self.config.llm_model_name.lower()
            if "gemini" in mn:
                try:
                    from lightrag.llm.gemini import gemini_model_complete
                    actual_llm_func = gemini_model_complete
                except ImportError:
                    pass
            elif any(k in mn for k in ("gpt-", "o1-", "o3-", "chatgpt")):
                try:
                    from lightrag.llm.openai import openai_complete_if_cache
                    actual_llm_func = openai_complete_if_cache
                except ImportError:
                    pass
            elif any(k in mn for k in ("llama", "mistral", "phi", "qwen", "deepseek", "tinyllama")):
                # HuggingFace local models (Colab/Kaggle) — common model families
                try:
                    from lightrag.llm.hf import hf_model_complete
                    actual_llm_func = hf_model_complete
                except ImportError:
                    pass
            # Ollama detection: explicit provider prefix or OLLAMA_HOST set
            if not actual_llm_func and (
                mn.startswith("ollama/") or os.getenv("OLLAMA_HOST")
            ):
                try:
                    from lightrag.llm.ollama import ollama_model_complete
                    actual_llm_func = ollama_model_complete
                    # Strip the "ollama/" prefix if present
                    if self.config.llm_model_name.lower().startswith("ollama/"):
                        self.config.llm_model_name = self.config.llm_model_name[7:]
                except ImportError:
                    pass
            
        # Determine base_url for LightRAG. When using native Gemini SDK functions
        # (gemini_model_complete, gemini_embed), the google-genai client handles 
        # its own endpoint — passing the OpenAI-bridge base_url causes 404s.
        lightrag_base_url = self.config.llm_api_base
        if self.config.llm_model_name and "gemini" in self.config.llm_model_name.lower():
            # The generativelanguage.googleapis.com URL is for the OpenAI bridge,
            # not for the native google-genai SDK. Pass None to let the SDK
            # use its default endpoint.
            if lightrag_base_url and "generativelanguage.googleapis.com" in lightrag_base_url:
                lightrag_base_url = None

        if actual_llm_func:
            lightrag_kwargs["llm_model_func"] = actual_llm_func
            # Pass API key, base URL, and model name via model_kwargs
            lightrag_kwargs["llm_model_kwargs"] = {
                "api_key": self.config.llm_api_key,
                "base_url": lightrag_base_url,
                "model_name": self.config.llm_model_name,
            }
            
        # 2. Resolve Embedding Function
        actual_embed_func = self.embedding_func
        if not actual_embed_func and self.config.embedding_model_name:
            en = self.config.embedding_model_name.lower()
            if "gemini" in en:
                try:
                    from lightrag.llm.gemini import gemini_embed
                    actual_embed_func = gemini_embed
                except ImportError:
                    pass
            elif any(k in en for k in ("text-embedding", "openai")):
                try:
                    from lightrag.llm.openai import openai_embed
                    actual_embed_func = openai_embed
                except ImportError:
                    pass
            # HuggingFace embedding models — set up via hf_embed with
            # user-provided tokenizer & model (initialized lazily)
            # Users pass model name like "BAAI/bge-small-en-v1.5"
            elif not actual_embed_func:
                try:
                    from lightrag.llm.hf import hf_embed
                    actual_embed_func = hf_embed
                except ImportError:
                    pass
        
        if actual_embed_func:
            from functools import partial
            from lightrag.utils import EmbeddingFunc
            from dataclasses import replace
            
            # If it's an EmbeddingFunc (common in LightRAG), we must use replace()
            # to keep its metadata (dim, model_name) while updating the partial func
            if isinstance(actual_embed_func, EmbeddingFunc):
                p_func = partial(
                    actual_embed_func.func,
                    api_key=self.config.llm_api_key,
                    base_url=lightrag_base_url,
                    model=self.config.embedding_model_name,
                )
                lightrag_kwargs["embedding_func"] = replace(
                    actual_embed_func, 
                    func=p_func,
                    embedding_dim=self.config.embedding_dim,
                    send_dimensions=True
                )
            else:
                # Fallback for plain functions
                lightrag_kwargs["embedding_func"] = partial(
                    actual_embed_func,
                    api_key=self.config.llm_api_key,
                    base_url=lightrag_base_url,
                    model=self.config.embedding_model_name,
                )

        # 3. Automatically use Supabase if credentials are provided
        if self.config.supabase_url and self.config.supabase_key:
            lightrag_kwargs["vector_storage"] = "SupabaseVectorStorage"
            
        self.lightrag = LightRAG(**lightrag_kwargs)
        await self.lightrag.initialize_storages()
        
        self._initialized = True
        
    async def close(self) -> None:
        """Close all connections and cleanup resources."""
        if self.graphiti:
            await self.graphiti.close()
            self.graphiti = None
            
        if self.lightrag:
            await self.lightrag.finalize_storages()
            self.lightrag = None
            
        self._initialized = False
        
    async def __aenter__(self):
        """Async context manager entry."""
        await self.initialize()
        return self
        
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """Async context manager exit."""
        await self.close()
        
    def _ensure_initialized(self) -> None:
        """Raise error if not initialized."""
        if not self._initialized:
            raise RuntimeError(
                "Helix not initialized. Call await helix.initialize() first, "
                "or use async context manager: async with Helix() as helix:"
            )
    
    async def insert(
        self,
        document: str,
        doc_id: str | None = None,
        source_description: str = "document",
        reference_time: datetime | None = None,
        group_id: str = "default",
        **kwargs,
    ) -> dict[str, Any]:
        """
        Insert document with temporal tracking.
        
        Uses Graphiti for Knowledge Graph production (entity/relation extraction)
        and LightRAG for vector indexing.
        
        Args:
            document: Text content to insert.
            doc_id: Optional document ID. Auto-generated if not provided.
            source_description: Description of document source.
            reference_time: When the facts in document were true (event time).
                           If None, uses current time.
            group_id: Group identifier for organizing related documents.
            **kwargs: Additional arguments passed to LightRAG.
            
        Returns:
            Dict with insert results including extracted entities and relations.
        """
        self._ensure_initialized()
        
        from graphiti_core.nodes import EpisodeType
        
        if reference_time is None:
            reference_time = datetime.now(timezone.utc)
            
        if doc_id is None:
            import hashlib
            doc_id = hashlib.md5(document.encode()).hexdigest()[:16]
        
        # Step 1: Add to Graphiti for temporal KG production
        # Graphiti handles:
        # - Entity extraction
        # - Entity resolution (deduplication)
        # - Edge invalidation (conflict resolution)
        # - Bi-temporal tracking (t_valid, t_created)
        graphiti_result = await self.graphiti.add_episode(
            name=f"doc_{doc_id}",
            episode_body=document,
            source=EpisodeType.text,
            source_description=source_description,
            reference_time=reference_time,
            group_id=group_id,
        )
        
        # Step 2: Add to LightRAG for vector indexing
        # LightRAG handles:
        # - Chunking
        # - Embedding generation
        # - Vector storage
        await self.lightrag.ainsert(document, **kwargs)
        
        return {
            "doc_id": doc_id,
            "episode_uuid": str(graphiti_result.episode.uuid),
            "entities_extracted": len(graphiti_result.nodes),
            "relations_extracted": len(graphiti_result.edges),
            "entity_names": [node.name for node in graphiti_result.nodes],
        }
    
    async def insert_batch(
        self,
        documents: list[str],
        doc_ids: list[str] | None = None,
        source_description: str = "batch_document",
        reference_times: list[datetime] | None = None,
        group_id: str = "default",
        **kwargs,
    ) -> list[dict[str, Any]]:
        """
        Insert multiple documents in batch.
        
        Args:
            documents: List of text documents.
            doc_ids: Optional list of document IDs.
            source_description: Description of document source.
            reference_times: List of reference times for each document.
            group_id: Group identifier.
            **kwargs: Additional arguments.
            
        Returns:
            List of insert results.
        """
        self._ensure_initialized()
        
        results = []
        for i, doc in enumerate(documents):
            doc_id = doc_ids[i] if doc_ids else None
            ref_time = reference_times[i] if reference_times else None
            
            result = await self.insert(
                document=doc,
                doc_id=doc_id,
                source_description=source_description,
                reference_time=ref_time,
                group_id=group_id,
                **kwargs,
            )
            results.append(result)
            
        return results
    
    async def query(
        self,
        query_text: str,
        mode: Literal["local", "global", "hybrid", "naive", "mix"] = "hybrid",
        top_k: int = 10,
        valid_at: datetime | None = None,
        include_temporal_context: bool = True,
        **kwargs,
    ) -> dict[str, Any]:
        """
        Query with temporal awareness.
        
        Combines Graphiti's temporal graph search with LightRAG's 
        dual-level retrieval for comprehensive results.
        
        Args:
            query_text: The query string.
            mode: Retrieval mode (local, global, hybrid, naive, mix).
            top_k: Number of results to retrieve.
            valid_at: Point-in-time for temporal queries.
                     If None, queries current state.
            include_temporal_context: Whether to include temporal metadata.
            **kwargs: Additional arguments passed to LightRAG.
            
        Returns:
            Dict with answer, context, and temporal metadata.
        """
        self._ensure_initialized()
        
        # Step 1: Search Graphiti for temporal graph context
        graphiti_results = await self.graphiti.search(
            query=query_text,
            num_results=top_k,
        )
        
        # Build temporal context from Graphiti results
        temporal_context = []
        for result in graphiti_results:
            temporal_context.append({
                "fact": result.fact if hasattr(result, 'fact') else str(result),
                "score": result.score if hasattr(result, 'score') else 0.0,
            })
        
        # Step 2: Query LightRAG for vector-based retrieval
        from lightrag.base import QueryParam
        lightrag_result = await self.lightrag.aquery(
            query_text,
            param=QueryParam(mode=mode, top_k=top_k, **kwargs)
        )
        
        result = {
            "answer": lightrag_result,
            "query": query_text,
            "mode": mode,
        }
        
        if include_temporal_context:
            result["temporal_context"] = temporal_context
            result["valid_at"] = valid_at.isoformat() if valid_at else "current"
            
        return result
    
    async def get_knowledge_graph(
        self,
        entity_name: str = "*",
        max_depth: int = 3,
        max_nodes: int = 100,
    ) -> dict[str, Any]:
        """
        Get knowledge graph visualization data.
        
        Args:
            entity_name: Starting entity name, or "*" for all.
            max_depth: Maximum traversal depth.
            max_nodes: Maximum nodes to return.
            
        Returns:
            Dict with nodes and edges for visualization.
        """
        self._ensure_initialized()
        
        # Use LightRAG's graph storage for this
        kg = await self.lightrag.get_knowledge_graph(
            node_label=entity_name,
            max_depth=max_depth,
            max_nodes=max_nodes,
        )
        
        return {
            "nodes": [
                {"id": n.id, "label": n.label, "properties": n.properties}
                for n in kg.nodes
            ] if hasattr(kg, 'nodes') else [],
            "edges": [
                {"source": e.source, "target": e.target, "label": e.label}
                for e in kg.edges
            ] if hasattr(kg, 'edges') else [],
            "is_truncated": kg.is_truncated if hasattr(kg, 'is_truncated') else False,
        }
