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

load_dotenv()

# ChromaDB configuration
CHROMA_DB_PATH = os.getenv("CHROMA_DB_PATH", "../chroma_db")
CHROMA_COLLECTION_NAME = os.getenv("CHROMA_COLLECTION_NAME", "generic_collection")

# Similarity parameters
SIMILARITY_TOP_K = int(os.getenv("SIMILARITY_TOP_K", "3"))

# Custom prompt template for query
CUSTOM_PROMPT_TEMPLATE = PromptTemplate(
    template="""
    System Instructions: You are KAFU Assistant for Kaimosi Friends University.
    ---------------------
    Context information:
    ---------------------
    {context_str}
    ---------------------
    Given the context information, answer the query below.

    Query: {query_str}
    Answer:
    """,
    prompt_type="query",
)

# Initialize ChromaDB
chroma_client = chromadb.PersistentClient(path=CHROMA_DB_PATH)
chroma_collection = chroma_client.get_or_create_collection(name=CHROMA_COLLECTION_NAME)
vector_store = ChromaVectorStore(chroma_collection=chroma_collection)

# 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  # Adjust as needed
