import streamlit as st
import os
from tempfile import NamedTemporaryFile
from langchain.document_loaders import PyPDFLoader
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import Chroma
from transformers import AutoModelForQuestionAnswering, AutoTokenizer, pipeline

# Function to save the uploaded PDF to a temporary file
def save_uploaded_file(uploaded_file):
    with NamedTemporaryFile(delete=False, suffix=".pdf") as temp_file:
        temp_file.write(uploaded_file.read())
        return temp_file.name

# Function to get answers from the PDF
def get_answer(question, db, model, tokenizer):
    doc = db.similarity_search(question, k=4)
    context = doc[0].page_content + doc[1].page_content + doc[2].page_content + doc[3].page_content

    # Load the model & tokenizer for question-answering
    model_name = "deepset/roberta-base-squad2"
    model = AutoModelForQuestionAnswering.from_pretrained(model_name)
    tokenizer = AutoTokenizer.from_pretrained(model_name)

    # Create a question-answering pipeline
    nlp = pipeline("question-answering", model=model, tokenizer=tokenizer)

    # Prepare the input
    QA_input = {
        "question": question,
        "context": context,
    }

    # Get the answer
    result = nlp(**QA_input)

    return result["answer"]

# Streamlit UI
st.title("PDF Question Answering App")
uploaded_file = st.file_uploader("Upload a PDF file", type=["pdf"])
if uploaded_file is not None:
    # Save the uploaded file to a temporary location
    temp_file_path = save_uploaded_file(uploaded_file)

    # Load the PDF document using PyPDFLoader
    loader = PyPDFLoader(temp_file_path)
    pages = loader.load_and_split()

    # Initialize embeddings and Chroma
    embed = HuggingFaceEmbeddings()
    db = Chroma.from_documents(pages, embed)

    # Load the model & tokenizer for question-answering
    model_name = "deepset/roberta-base-squad2"
    model = AutoModelForQuestionAnswering.from_pretrained(model_name)
    tokenizer = AutoTokenizer.from_pretrained(model_name)

    # Initializations
    conversation = []
    st.write("Ask your questions, and I'll provide answers:")
    
    # Continuous question-answering loop
    while True:
        question = st.text_input("Enter your question:")
        if st.button("Get Answer"):
            answer = get_answer(question, db, model, tokenizer)
            st.write("Answer:")
            st.write(answer)
            conversation.append({"question": question, "answer": answer})
        
        # Add an option to end the conversation
        if st.button("End Conversation"):
            break

    # Display the conversation history
    st.write("Conversation History:")
    for entry in conversation:
        st.write(f"Q: {entry['question']}")
        st.write(f"A: {entry['answer']}")

    # Cleanup: Delete the temporary file
    os.remove(temp_file_path)