import os
import streamlit as st
from langchain_google_genai import GoogleGenerativeAIEmbeddings
import google.generativeai as genai
from langchain_community.vectorstores import FAISS
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain.chains.question_answering import load_qa_chain
from langchain.prompts import PromptTemplate
from dotenv import load_dotenv

load_dotenv()
genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))

@st.cache_resource
def load_models_and_embeddings():
    embeddings = GoogleGenerativeAIEmbeddings(model="models/embedding-001")
    conversational_model = ChatGoogleGenerativeAI(model="gemini-2.5-flash", temperature=0.7)
    return embeddings, conversational_model

def get_conversational_chain(conversational_model):
    prompt_template = """
    You are a helpful assistant tasked with extracting accurate answers **only from the given context**.

    If the question is about matching (e.g., "Who is referred to as X?"), select the correct match from the context.

    If the correct answer is **not present** in the context, respond exactly with:
    "উত্তর প্রসঙ্গে নেই" (The answer is not in the context.)

    ---

    প্রসঙ্গ (Context):
    {context}

    প্রশ্ন (Question):
    {question}

    উত্তর (Answer):
    """
    prompt = PromptTemplate(template=prompt_template, input_variables=["context", "question"])
    chain = load_qa_chain(conversational_model, chain_type="stuff", prompt=prompt)
    return chain

def user_input(user_question, embeddings, conversational_model):
    try:
        with st.spinner("Generating answer..."):
            new_db = FAISS.load_local("faiss_index", embeddings, allow_dangerous_deserialization=True)
            docs_and_scores = new_db.similarity_search_with_score(user_question, k=100)
            docs = [doc for doc, score in docs_and_scores]
            chain = get_conversational_chain(conversational_model)
            response = chain({"input_documents": docs, "question": user_question})
            st.write("### Answer:")
            st.write(response["output_text"])
    except Exception as e:
        st.error(f"An error occurred: {e}")

def main():
    st.set_page_config(page_title="AEA")
    st.header("AEA")
    embeddings, conversational_model = load_models_and_embeddings()
    
    with st.form(key="qa_form"):
        user_question = st.text_input("Ask AEA")
        submitted = st.form_submit_button("Submit")

    if submitted and user_question:
        user_input(user_question, embeddings, conversational_model)

if __name__ == "__main__":
    main()