"""
Helix: Temporal GraphRAG — LightRAG + Graphiti Integration.

Architecture:
  LightRAG  = Sole orchestrator (chunking, embedding, retrieval, LLM gen)
  Graphiti  = KG intelligence (entity resolution, dedup, temporal tracking)
  Adapter   = GraphitiStorage (thin SDK bridge registered as LightRAG plugin)

All data flows through ONE pipeline: LightRAG controls everything,
Graphiti is plugged in via GraphitiStorage.
"""

__version__ = "0.2.0"
__author__ = "Yash Nuhash"

import os
import logging

# ---------------------------------------------------------------------------
# Runtime registration of Helix storage backends into LightRAG's registry.
# ---------------------------------------------------------------------------
try:
    from lightrag.kg import STORAGES, STORAGE_IMPLEMENTATIONS

    # --- Graph Storage: GraphitiStorage ---
    if "GraphitiStorage" not in STORAGES:
        STORAGES["GraphitiStorage"] = "helix.storage.graphiti_impl"
    _graph_impls = STORAGE_IMPLEMENTATIONS.get("GRAPH_STORAGE", {}).get("implementations", [])
    if "GraphitiStorage" not in _graph_impls:
        _graph_impls.append("GraphitiStorage")

    # --- Vector Storage: SupabaseVectorStorage ---
    if "SupabaseVectorStorage" not in STORAGES:
        STORAGES["SupabaseVectorStorage"] = "helix.storage.supabase_impl"
    _vec_impls = STORAGE_IMPLEMENTATIONS.get("VECTOR_STORAGE", {}).get("implementations", [])
    if "SupabaseVectorStorage" not in _vec_impls:
        _vec_impls.append("SupabaseVectorStorage")

except ImportError:
    pass  # lightrag not installed yet; will fail later with a clear message

# ---------------------------------------------------------------------------
# Global Gemini patch — ensure ALL GeminiClients use the configured model
# and don't fall back to models unavailable on restricted API keys.
# ---------------------------------------------------------------------------
try:
    import graphiti_core.llm_client.gemini_client as gc_mod
    from graphiti_core.llm_client.gemini_client import GeminiClient

    if not hasattr(gc_mod, "_helix_patched"):
        original_init = GeminiClient.__init__

        def patched_init(self, *args, **kwargs):
            original_init(self, *args, **kwargs)
            if hasattr(self, "config") and self.config and getattr(self.config, "model", None):
                self.small_model = self.config.model
            elif hasattr(self, "model") and self.model:
                self.small_model = self.model

        def patched_get_model(self, model_size):
            return self.model

        GeminiClient.__init__ = patched_init
        GeminiClient._get_model_for_size = patched_get_model
        gc_mod._helix_patched = True

except ImportError:
    pass
except Exception:
    pass

from helix.core import Helix, HelixConfig

__all__ = [
    "__version__",
    "__author__",
    "Helix",
    "HelixConfig",
]
