from __future__ import annotations

from typing import Any

from fastapi_modules.fastapi_leaudit.rag_engine.config import RAG_CONFIG

_instance: Any | None = None


def init_chroma() -> Any:
    global _instance
    if _instance is not None:
        return _instance

    import chromadb  # lazy import to avoid hard failure before feature is enabled
    import chromadb.config

    host = RAG_CONFIG["CHROMA_HOST"]
    if host:
        token = RAG_CONFIG.get("CHROMA_TOKEN", "")
        header = RAG_CONFIG.get("CHROMA_AUTH_HEADER", "X-Chroma-Token")
        settings = (
            chromadb.config.Settings(
                chroma_client_auth_provider="chromadb.auth.token_authn.TokenAuthClientProvider",
                chroma_client_auth_credentials=token,
                chroma_auth_token_transport_header=header,
            )
            if token
            else chromadb.config.Settings()
        )
        _instance = chromadb.HttpClient(host=host, port=RAG_CONFIG["CHROMA_PORT"], settings=settings)
    else:
        _instance = chromadb.PersistentClient(path=RAG_CONFIG["CHROMA_PERSIST_DIR"])
    return _instance


def get_chroma() -> Any:
    if _instance is None:
        return init_chroma()
    return _instance