from langchain_community.vectorstores import FAISS
from langchain_core.documents import Document
from langchain_nvidia_ai_endpoints import NVIDIAEmbeddings
import os 
from google import genai
from google.genai import types

# Set up the Gemini API key
import os

def index_text():

    os.environ["NVIDIA_API_KEY"] =os.getenv("NVIDIA_API_KEY")

    
    nvidia_embeddings = NVIDIAEmbeddings(
        model="nvidia/llama-3.2-nv-embedqa-1b-v2",
        truncate="NONE"
    )
    vectorstore = FAISS.load_local("nvidia_faiss_index", embeddings=nvidia_embeddings,allow_dangerous_deserialization=True)
    return vectorstore


def answer_query(query, history,vectorstore):

    os.environ["GEMINI_API_KEY"] = os.getenv("GEMINI_API_KEY")
    client = genai.Client()

    RAG_TEMPLATE = """
#CONTEXT:
{context}
Use the provided context to answer the user query.
"""
    retriever = vectorstore.as_retriever()
    search_results = retriever.invoke(query, k=2)
    context = " ".join([doc.page_content for doc in search_results])
    prompt = RAG_TEMPLATE.format(context=context, query=query)


    gemini_history = []
    for msg in history:
        # The Gemini API uses 'model' for the assistant's role
        
        role = 'model' if msg['role'] == 'assistant' else 'user'
        gemini_history.append(
            types.Content(role=role, parts=[types.Part(text=msg['content'])])
        )

    chat = client.chats.create(
        model="gemini-2.0-flash",
        history=gemini_history,
        config=types.GenerateContentConfig(
            system_instruction=prompt)
    )

    response=chat.send_message(message=query)
    return response.text