"""
Helix: Thin configuration layer over LightRAG.

In the ideal architecture LightRAG is the **sole orchestrator**.
Helix's job is purely to:
  1. Read environment / HelixConfig
  2. Wire up GraphitiStorage (the adapter) as LightRAG's graph backend
  3. Optionally wire up SupabaseVectorStorage as the vector backend
  4. Expose convenience helpers (insert / query / close) that delegate
     100% to LightRAG — no parallel Graphiti calls.

All Graphiti interaction happens inside GraphitiStorage, which is loaded
by LightRAG as a normal storage plugin.
"""

from __future__ import annotations

import os
import logging
from dataclasses import dataclass, field
from typing import Any, Callable, Literal

from dotenv import load_dotenv

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

logger = logging.getLogger("helix.core")


# ────────────────────────────────────────────────────────────────────
#  Configuration
# ────────────────────────────────────────────────────────────────────

@dataclass
class HelixConfig:
    """Configuration for a Helix instance.

    All values fall back to environment variables so that a bare
    ``Helix()`` works when the user has a ``.env`` file.
    """

    # Working directory for LightRAG artefacts
    working_dir: str = field(default="./helix_storage")

    # Neo4j — used by GraphitiStorage
    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 — optional 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
    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
    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"))
    )


# ────────────────────────────────────────────────────────────────────
#  Helix — LightRAG wrapper
# ────────────────────────────────────────────────────────────────────

class Helix:
    """LightRAG + Graphiti, unified through the storage-plugin model.

    All data flows through **one** pipeline::

        Client -> Helix.insert / Helix.query
                  |-> LightRAG.ainsert / aquery
                      |-> Chunker / NLP  (operate.py)
                      |-> Embedding svc
                      |-> Vector Storage
                      |-> GraphitiStorage  (adapter)
                      |     |-> Graphiti Core
                      |          |-> Neo4j Driver
                      |-> LLM Generation

    Helix **never** calls ``graphiti.search()`` or ``graphiti.add_episode()``
    directly.  Those calls live inside ``GraphitiStorage``.
    """

    def __init__(
        self,
        config: HelixConfig | None = None,
        working_dir: str | None = None,
        llm_model_func: Callable | None = None,
        embedding_func: Callable | None = None,
    ):
        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

        # Populated in initialize()
        self.lightrag = None
        self._initialized = False

    # ----------------------------------------------------------------
    #  Lifecycle
    # ----------------------------------------------------------------

    async def initialize(self) -> None:
        """Build LightRAG with GraphitiStorage plugged in."""
        if self._initialized:
            return

        try:
            from lightrag import LightRAG
        except ImportError:
            raise ImportError(
                "lightrag is required for Helix. "
                "Install with: pip install lightrag-hku"
            )

        # Propagate Neo4j env-vars so GraphitiStorage can read them
        os.environ.setdefault("NEO4J_URI", self.config.neo4j_uri)
        os.environ.setdefault("NEO4J_USERNAME", self.config.neo4j_user)
        os.environ.setdefault("NEO4J_PASSWORD", self.config.neo4j_password)
        os.environ.setdefault("NEO4J_DATABASE", self.config.neo4j_database)

        # -- kwargs for GraphitiStorage (passed via graph_storage_cls_kwargs) --
        graph_kwargs: dict[str, Any] = {
            "neo4j_uri": self.config.neo4j_uri,
            "neo4j_user": self.config.neo4j_user,
            "neo4j_password": self.config.neo4j_password,
            "neo4j_database": self.config.neo4j_database,
            "llm_model_name": self.config.llm_model_name,
            "llm_api_key": self.config.llm_api_key,
            "embedding_dim": self.config.embedding_dim,
        }

        # -- LightRAG init kwargs --
        lightrag_kwargs: dict[str, Any] = {
            "working_dir": self.config.working_dir,
        }

        # Check if installed LightRAG supports graph_storage_cls_kwargs
        # (present in repo's local lightrag/ but may be absent in pip lightrag-hku)
        import inspect
        _lr_params = inspect.signature(LightRAG).parameters
        _supports_graph_cls_kwargs = "graph_storage_cls_kwargs" in _lr_params

        if _supports_graph_cls_kwargs:
            lightrag_kwargs["graph_storage"] = "GraphitiStorage"
            lightrag_kwargs["graph_storage_cls_kwargs"] = graph_kwargs
            logger.info("Using GraphitiStorage with graph_storage_cls_kwargs")
        else:
            # Fallback: store graph kwargs in env-vars for GraphitiStorage
            # to read, and skip the unsupported kwarg entirely.
            logger.warning(
                "Installed LightRAG does not support graph_storage_cls_kwargs. "
                "Using default graph storage with Neo4j env-vars."
            )

        # -- Resolve LLM function --
        actual_llm_func = self._resolve_llm_func()
        lightrag_base_url = self._resolve_base_url()

        if actual_llm_func:
            lightrag_kwargs["llm_model_func"] = actual_llm_func
            lightrag_kwargs["llm_model_kwargs"] = {
                "api_key": self.config.llm_api_key,
                "base_url": lightrag_base_url,
                "model_name": self.config.llm_model_name,
            }
            lightrag_kwargs["llm_model_name"] = self.config.llm_model_name

        # -- Resolve Embedding function --
        actual_embed_func = self._resolve_embed_func()
        if actual_embed_func:
            from functools import partial
            from lightrag.utils import EmbeddingFunc
            from dataclasses import replace as dc_replace

            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"] = dc_replace(
                    actual_embed_func,
                    func=p_func,
                    embedding_dim=self.config.embedding_dim,
                    send_dimensions=True,
                )
            else:
                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,
                )

        # -- Supabase vector storage (optional) --
        if self.config.supabase_url and self.config.supabase_key:
            lightrag_kwargs["vector_storage"] = "SupabaseVectorStorage"

        # -- Create & initialise LightRAG --
        try:
            self.lightrag = LightRAG(**lightrag_kwargs)
        except TypeError as e:
            # Handle any remaining unexpected kwargs for pip-installed versions
            unsupported = str(e)
            logger.warning(f"LightRAG init issue: {unsupported}")
            # Strip unsupported kwargs and retry
            for kwarg_name in list(lightrag_kwargs.keys()):
                if kwarg_name in unsupported:
                    logger.info(f"  Removing unsupported kwarg: {kwarg_name}")
                    del lightrag_kwargs[kwarg_name]
            self.lightrag = LightRAG(**lightrag_kwargs)

        await self.lightrag.initialize_storages()

        self._initialized = True
        logger.info("Helix initialised — LightRAG is the sole orchestrator.")

    async def close(self) -> None:
        """Tear down LightRAG (which tears down GraphitiStorage -> Graphiti)."""
        if self.lightrag:
            await self.lightrag.finalize_storages()
            self.lightrag = None
        self._initialized = False

    async def __aenter__(self):
        await self.initialize()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await self.close()

    # ----------------------------------------------------------------
    #  Insert — SINGLE pipeline through LightRAG
    # ----------------------------------------------------------------

    async def insert(
        self,
        document: str | list[str],
        doc_id: str | None = None,
        ids: list[str] | None = None,
        file_paths: str | list[str] | None = None,
        **kwargs,
    ) -> str:
        """Insert document(s) into the RAG pipeline.

        Everything flows through ``LightRAG.ainsert()`` which:
        1. Chunks the document            (operate.py)
        2. Extracts entities / relations   (operate.py -> LLM)
        3. Writes to GraphitiStorage       (adapter -> Graphiti SDK)
        4. Embeds chunks into VectorDB     (embedding_func)

        After entity/relation writes complete, ``GraphitiStorage.index_done_callback()``
        fires ``graphiti.add_episode()`` so Graphiti Core can run entity resolution,
        deduplication, and temporal tracking.

        Returns:
            str: tracking ID for monitoring processing status.
        """
        self._ensure_initialized()
        if ids is None and doc_id is not None:
            ids = [doc_id]
        return await self.lightrag.ainsert(
            document, ids=ids, file_paths=file_paths, **kwargs
        )

    async def insert_batch(
        self,
        documents: list[str],
        doc_ids: list[str] | None = None,
        file_paths: list[str] | None = None,
        **kwargs,
    ) -> str:
        """Insert multiple documents — single pipeline."""
        self._ensure_initialized()
        return await self.lightrag.ainsert(
            documents, ids=doc_ids, file_paths=file_paths, **kwargs
        )

    # ----------------------------------------------------------------
    #  Query — SINGLE pipeline through LightRAG
    # ----------------------------------------------------------------

    async def query(
        self,
        query_text: str,
        mode: Literal["local", "global", "hybrid", "naive", "mix"] = "hybrid",
        top_k: int = 10,
        **kwargs,
    ) -> str:
        """Query the RAG pipeline.

        LightRAG's retriever reads from GraphitiStorage (which queries
        Graphiti's graph data in Neo4j) and from the VectorDB, then
        assembles a prompt and calls the LLM.

        Returns:
            str: The LLM-generated answer.
        """
        self._ensure_initialized()
        from lightrag.base import QueryParam
        return await self.lightrag.aquery(
            query_text,
            param=QueryParam(mode=mode, top_k=top_k, **kwargs),
        )

    # ----------------------------------------------------------------
    #  Convenience — access the underlying storage for advanced ops
    # ----------------------------------------------------------------

    @property
    def graph_storage(self):
        """Return the GraphitiStorage instance (adapter layer)."""
        self._ensure_initialized()
        return self.lightrag.chunk_entity_relation_graph

    @property
    def graphiti(self):
        """Return the Graphiti instance managed by GraphitiStorage.

        Useful for temporal queries, multi-hop, hallucination detection
        that need direct access to Graphiti's ``search()`` method.
        """
        gs = self.graph_storage
        return getattr(gs, "graphiti", None)

    async def get_knowledge_graph(
        self,
        entity_name: str = "*",
        max_depth: int = 3,
        max_nodes: int = 100,
    ) -> dict[str, Any]:
        """Get KG data for visualisation."""
        self._ensure_initialized()
        kg = await self.lightrag.chunk_entity_relation_graph.get_knowledge_graph(
            node_label=entity_name,
            max_depth=max_depth,
            max_nodes=max_nodes,
        )
        return {
            "nodes": [
                {"id": n.get("id", ""), "label": n.get("label", ""), "properties": n}
                for n in (kg.nodes if hasattr(kg, "nodes") else [])
            ],
            "edges": [
                {
                    "source": e.get("source", ""),
                    "target": e.get("target", ""),
                    "label": e.get("label", ""),
                }
                for e in (kg.edges if hasattr(kg, "edges") else [])
            ],
            "is_truncated": getattr(kg, "is_truncated", False),
        }

    # ----------------------------------------------------------------
    #  Private helpers
    # ----------------------------------------------------------------

    def _ensure_initialized(self) -> None:
        if not self._initialized:
            raise RuntimeError(
                "Helix not initialised. Call `await helix.initialize()` first, "
                "or use `async with Helix() as helix:`"
            )

    def _resolve_llm_func(self) -> Callable | None:
        if self.llm_model_func:
            return self.llm_model_func
        mn = (self.config.llm_model_name or "").lower()
        if not mn:
            return None

        if "gemini" in mn:
            try:
                from lightrag.llm.gemini import gemini_model_complete
                return 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
                return openai_complete_if_cache
            except ImportError:
                pass
        elif any(k in mn for k in ("llama", "mistral", "phi", "qwen", "deepseek", "tinyllama")):
            try:
                from lightrag.llm.hf import hf_model_complete
                return hf_model_complete
            except ImportError:
                pass

        # Ollama fallback
        if mn.startswith("ollama/") or os.getenv("OLLAMA_HOST"):
            try:
                from lightrag.llm.ollama import ollama_model_complete
                if mn.startswith("ollama/"):
                    self.config.llm_model_name = self.config.llm_model_name[7:]
                return ollama_model_complete
            except ImportError:
                pass

        return None

    def _resolve_embed_func(self):
        if self.embedding_func:
            return self.embedding_func
        en = (self.config.embedding_model_name or "").lower()
        if not en:
            return None

        if "gemini" in en:
            try:
                from lightrag.llm.gemini import gemini_embed
                return gemini_embed
            except ImportError:
                pass
        elif any(k in en for k in ("text-embedding", "openai")):
            try:
                from lightrag.llm.openai import openai_embed
                return openai_embed
            except ImportError:
                pass
        else:
            try:
                from lightrag.llm.hf import hf_embed
                return hf_embed
            except ImportError:
                pass
        return None

    def _resolve_base_url(self) -> str | None:
        url = self.config.llm_api_base
        if self.config.llm_model_name and "gemini" in self.config.llm_model_name.lower():
            if url and "generativelanguage.googleapis.com" in url:
                return None
        return url or None
