import os
import gradio as gr
from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
import google.generativeai as genai
import tempfile
import shutil

# Local vector store path
INDEX_PATH = "faiss_index"

# Load or initialize vector DB
def load_vectorstore(pdf_path):
    loader = PyPDFLoader(pdf_path)
    docs = loader.load()

    splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=100)
    chunks = splitter.split_documents(docs)

    embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
    vectorstore = FAISS.from_documents(chunks, embeddings)

    # Save locally
    vectorstore.save_local(INDEX_PATH)
    return vectorstore

def load_existing_vectorstore():
    embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
    return FAISS.load_local(INDEX_PATH, embeddings)

# Gemini: Ask question using optimized context
def ask_question(api_key, question, pdf_file):
    if not api_key or not api_key.startswith("AI"):
        return "❌ Invalid Gemini API Key. Please enter a valid key starting with 'AI'."

    # Setup Gemini
    genai.configure(api_key=api_key)
    model = genai.GenerativeModel("models/gemini-1.5-flash")

    # Get uploaded file path
    pdf_path = pdf_file.name

    # Load or create vectorstore
    if os.path.exists(INDEX_PATH):
        vs = load_existing_vectorstore()
    else:
        vs = load_vectorstore(pdf_path)

    # Search for relevant chunks (context optimization)
    top_docs = vs.similarity_search(question, k=3)
    context = "\n\n".join([doc.page_content for doc in top_docs])

    # Construct the prompt
    prompt = f"""Use the context below to answer the question:\n\nContext:\n{context}\n\nQuestion: {question}"""

    try:
        response = model.generate_content(prompt)
        return response.text
    except Exception as e:
        return f"❌ Gemini Error: {str(e)}"

# Gradio UI
with gr.Blocks(title="RAG Q&A with Gemini") as demo:
    gr.Markdown("## 📄 Upload PDF | 🔍 Ask Questions | 🔌 Powered by Gemini")
    api_key = gr.Textbox(label="🔐 Gemini API Key", placeholder="Enter your Gemini API key")
    pdf_file = gr.File(label="📂 Upload PDF", file_types=[".pdf"])
    question = gr.Textbox(label="❓ Your Question")
    answer = gr.Textbox(label="💬 Answer", lines=8)
    ask_btn = gr.Button("🚀 Ask Gemini")

    ask_btn.click(fn=ask_question, inputs=[api_key, question, pdf_file], outputs=answer)

# Run the app
if __name__ == "__main__":
    demo.launch()