import os
from dotenv import load_dotenv
from llama_index.vector_stores.chroma import ChromaVectorStore
from llama_index.core import VectorStoreIndex, StorageContext, Settings
from llama_index.llms.gemini import Gemini
from llama_index.embeddings.gemini import GeminiEmbedding
import chromadb


def main():
    # Load environment variables
    load_dotenv()

    # Initialize Gemini LLM and embedding model
    llm = Gemini(api_key=os.environ["GEMINI_API_KEY"])
    embed_model = GeminiEmbedding(model_name="models/embedding-001")

    # Set global settings
    Settings.llm = llm
    Settings.embed_model = embed_model
    Settings.chunk_size = 1024

    # Initialize Chroma client and collection (already trained)
    chroma_client = chromadb.PersistentClient(path="../chroma_db")
    chroma_collection = chroma_client.get_or_create_collection(name="kafu")

    # Set up vector store and storage context
    vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
    storage_context = StorageContext.from_defaults(vector_store=vector_store)

    # Load index directly from vector store
    index = VectorStoreIndex.from_vector_store(vector_store)

    # Set up query engine
    query_engine = index.as_query_engine()

    # Run query
    query = "What is KAFU"
    response = query_engine.query(query)

    print(f"\nQuery: {query}\nResponse:\n{response}")


if __name__ == "__main__":
    main()
