import gradio as gr
from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.llms import HuggingFaceHub


# -----------------------------
# LOAD & SPLIT PDF
# -----------------------------
def process_pdf(file):
    loader = PyPDFLoader(file.name)
    pages = loader.load()

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

    docs = splitter.split_documents(pages)
    return docs


# -----------------------------
# CREATE VECTOR DB
# -----------------------------
def create_db(docs):
    embeddings = HuggingFaceEmbeddings(
        model_name="sentence-transformers/all-MiniLM-L6-v2"
    )

    db = FAISS.from_documents(docs, embeddings)
    return db


# -----------------------------
# QUESTION ANSWERING
# -----------------------------
def ask_question(file, question):
    if file is None:
        return "Please upload a PDF first."

    docs = process_pdf(file)
    db = create_db(docs)

    retrieved_docs = db.similarity_search(question, k=5)

    context = "\n\n".join([doc.page_content for doc in retrieved_docs])

    if not context.strip():
        return "No relevant information found in the document."

    prompt = f"""
You are a helpful assistant. Answer ONLY using the given context.

If the answer is not present, reply:
"Not found in report"

Context:
{context}

Question:
{question}

Answer:
"""

    llm = HuggingFaceHub(
        repo_id="google/flan-t5-base",
        model_kwargs={"temperature": 0.3, "max_length": 256}
    )

    result = llm.invoke(prompt)
    return result


# -----------------------------
# GRADIO UI
# -----------------------------
with gr.Blocks() as app:
    gr.Markdown("# 🧠 Medical Report Q&A (RAG)")

    file = gr.File(label="Upload PDF")
    question = gr.Textbox(label="Ask your question")

    btn = gr.Button("Get Answer")
    output = gr.Textbox(label="Answer")

    btn.click(fn=ask_question, inputs=[file, question], outputs=output)

app.launch()