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

Helix combines LightRAG's proven dual-level retrieval with Graphiti's
bi-temporal Knowledge Graph to create a next-generation RAG system.
"""

__version__ = "0.1.8"
__author__ = "Yash Nuhash"

import os
import logging

# ---------------------------------------------------------------------------
# Runtime registration of Helix storage backends into LightRAG's registry.
#
# This allows `pip install helix-rag lightrag-hku` to work without forking
# or publishing a custom LightRAG package.  We simply inject our entries
# into the two dicts that LightRAG's _get_storage_class() consults at
# import-time, before any LightRAG instance is created.
# ---------------------------------------------------------------------------
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 Monkeypatch for Gemini — ensure ALL GeminiClients use the correct
# model and don't fall back to models unavailable on restricted API keys.
# Must be applied before any Graphiti instances are created.
# ---------------------------------------------------------------------------
try:
    print("DEBUG: Helix package initializing... applying Gemini mitigation.")
    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
        print("DEBUG: GeminiClient successfully patched.")

except ImportError:
    # graphiti-core may not be installed yet
    pass
except Exception as patch_err:
    print(f"DEBUG: Failed to apply Gemini patch: {patch_err}")

from helix.core import Helix, HelixConfig

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