import gradio as gr
from langchain.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import FAISS
from langchain.chains import RetrievalQA
from langchain.llms import HuggingFacePipeline

from transformers import pipeline
import os

# Global variables
vectorstore = None
qa_chain = None

# Load embedding model once
embeddings = HuggingFaceEmbeddings(
    model_name="sentence-transformers/all-MiniLM-L6-v2"
)

# Load LLM once
pipe = pipeline(
    "text2text-generation",
    model="google/flan-t5-base",
    max_length=512
)

llm = HuggingFacePipeline(pipeline=pipe)

def process_pdf(pdf_file):
    global vectorstore, qa_chain
    
    if pdf_file is None:
        return "Please upload a PDF first."
    
    # Load PDF
    loader = PyPDFLoader(pdf_file.name)
    documents = loader.load()
    
    # Split text
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000,
        chunk_overlap=200
    )
    texts = splitter.split_documents(documents)
    
    # Create vector DB
    vectorstore = FAISS.from_documents(texts, embeddings)
    
    # Create QA chain
    qa_chain = RetrievalQA.from_chain_type(
        llm=llm,
        retriever=vectorstore.as_retriever(),
        return_source_documents=True
    )
    
    return "PDF processed successfully! You can now ask questions."

def ask_question(question):
    global qa_chain
    
    if qa_chain is None:
        return "Upload and process a PDF first."
    
    result = qa_chain(question)
    
    answer = result["result"]
    sources = "\n\n".join(
        [doc.page_content[:300] for doc in result["source_documents"]]
    )
    
    return f"Answer:\n{answer}\n\nSources:\n{sources}"

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("# 📄 PDF Question Answering System")
    gr.Markdown("Upload a PDF and ask questions about it.")
    
    pdf_input = gr.File(file_types=[".pdf"])
    process_btn = gr.Button("Process PDF")
    status = gr.Textbox(label="Status")
    
    question = gr.Textbox(label="Ask a question")
    ask_btn = gr.Button("Get Answer")
    
    output = gr.Textbox(label="Response", lines=15)
    
    process_btn.click(process_pdf, inputs=pdf_input, outputs=status)
    ask_btn.click(ask_question, inputs=question, outputs=output)

demo.launch()