import gradio as gr
import torch

from transformers import pipeline

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


# -----------------------------
# LOCAL CPU MODEL
# -----------------------------
pipe = pipeline(
    "text2text-generation",
    model="google/flan-t5-base",
    max_length=512,
    device=-1  # CPU
)

llm = HuggingFacePipeline(pipeline=pipe)


# -----------------------------
# VECTORSTORE BUILDER
# -----------------------------
def build_vectorstore_from_file(pdf_path):
    loader = PyPDFLoader(pdf_path)
    documents = loader.load()

    splitter = RecursiveCharacterTextSplitter(
        chunk_size=800,
        chunk_overlap=150
    )

    chunks = splitter.split_documents(documents)

    embeddings = HuggingFaceEmbeddings(
        model_name="sentence-transformers/all-MiniLM-L6-v2"
    )

    return FAISS.from_documents(chunks, embeddings)


# -----------------------------
# QA FUNCTION
# -----------------------------
def answer_question(pdf, question):
    try:
        if pdf is None:
            return "❌ Upload a PDF first"

        if not question.strip():
            return "❌ Ask a question"

        vectorstore = build_vectorstore_from_file(pdf)

        qa = RetrievalQA.from_chain_type(
            llm=llm,
            retriever=vectorstore.as_retriever(search_kwargs={"k": 3}),
            chain_type="stuff"
        )

        return qa.run(question)

    except Exception as e:
        return f"❌ Error: {str(e)}"


# -----------------------------
# GRADIO UI
# -----------------------------
with gr.Blocks() as demo:
    gr.Markdown("## 📄 PDF Question Answer Bot")

    pdf = gr.File(label="Upload PDF", file_types=[".pdf"])
    question = gr.Textbox(label="Question")
    answer = gr.Textbox(label="Answer", lines=6)

    btn = gr.Button("Ask")
    btn.click(answer_question, [pdf, question], answer)

demo.launch()