import streamlit as st
from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_huggingface import HuggingFaceEmbeddings
from langchain.vectorstores import Chroma
from langchain.prompts import ChatPromptTemplate
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_core.runnables import RunnablePassthrough
import tempfile


# --- Streamlit App Title ---
st.title("📄 Medical Education Visual Aid - PDF RAG Chatbot")

# --- File Uploader ---
uploaded_file = st.file_uploader("Upload a PDF document", type=["pdf"])

# Check if a file is uploaded
if uploaded_file:
    # Save uploaded file to a temporary path
    with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp_file:
        tmp_file.write(uploaded_file.read())
        temp_pdf_path = tmp_file.name

    # Load and process the PDF
    loader = PyPDFLoader(temp_pdf_path)
    pages = loader.load()

    splitter = RecursiveCharacterTextSplitter(
        chunk_size=100,
        chunk_overlap=20
    )
    split = splitter.split_documents(pages)

    # Create embeddings and vector store
    hfembedding = HuggingFaceEmbeddings()
    db = Chroma.from_documents(documents=split, embedding=hfembedding)

    # Prompt template
    template_string = """
    You are a smart medical assistant. Use the provided context to answer the question.
    context: {context}
    question: {question}
    """
    prompt_template = ChatPromptTemplate.from_template(template_string)

    # LLM initialization (replace with your key or env variable)
    GOOGLE_API_KEY = "AIzaSyDA1zZ006f4krgUDjhdju_RcwZNwNPOJwI"
    llm = ChatGoogleGenerativeAI(model="gemini-1.5-flash-latest", api_key=GOOGLE_API_KEY)

        # RAG chain
    retriever = db.as_retriever()
    chain = (
            {"context": retriever, "question": RunnablePassthrough()}
            | prompt_template
            | llm
        )

        # User query
    query = st.text_input("💬 Ask a question about the uploaded PDF:")
    if query:
          response = chain.invoke(query)
          st.subheader("Answer")
          st.write(response.content)