import gradio as gr
import os
from langchain.vectorstores import Chroma
from langchain_community.document_loaders import PyPDFLoader
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate
import fitz  # PyMuPDF für das Extrahieren von Text aus PDFs

# Funktion zum Extrahieren von Text aus einer PDF
def extract_text_from_pdf(pdf_path):
    doc = fitz.open(pdf_path)  # Öffnen der PDF-Datei
    text_pages = []
    
    # Durch alle Seiten der PDF iterieren und Text extrahieren
    for page_num in range(doc.page_count):
        page = doc.load_page(page_num)
        text = page.get_text("text")  # Extrahiert den Text als "plain text"
        text_pages.append(text)
    
    return text_pages

# Frage-Antwort-Funktion mit Langchain und Chroma
def process_pdf_and_query(pdf_path, question):
    # Extrahiere Text aus der PDF
    extracted_text = extract_text_from_pdf(pdf_path)
    
    # Dokumente für Langchain laden
    documents = [{"text": page_text} for page_text in extracted_text]

    # Embedding und Vektorstore vorbereiten
    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
    vectordb = Chroma.from_documents(documents, embeddings)

    retriever = vectordb.as_retriever()
    prompt_template = "Beantworte die folgende Frage basierend auf dem Dokument: {context}\nFrage: {question}\nAntwort:"
    prompt = PromptTemplate(input_variables=["context", "question"], template=prompt_template)

    # Erstellung der RetrievalQA-Kette
    qa_chain = RetrievalQA.from_chain_type(llm=None, retriever=retriever, chain_type_kwargs={"prompt": prompt})
    response = qa_chain.run(input_documents=documents, question=question)
    
    return response

# Gradio Antwortfunktion
def chatbot_response(pdf, question):
    # Speichern der hochgeladenen PDF-Datei
    pdf_path = "/tmp/uploaded_pdf.pdf"
    pdf_content = pdf.read()  # Hole den Inhalt der PDF als Byte-Stream
    
    # Speichern des Byte-Streams in einer Datei
    with open(pdf_path, "wb") as f:
        f.write(pdf_content)
    
    # Frage beantworten basierend auf der extrahierten PDF und der Frage
    answer = process_pdf_and_query(pdf_path, question)

    # Temporäre Datei löschen
    os.remove(pdf_path)
    
    return answer

# Gradio Interface erstellen
pdf_input = gr.File(label="PDF-Datei hochladen")
question_input = gr.Textbox(label="Frage eingeben")
response_output = gr.Textbox(label="Antwort")

# Gradio Interface starten
interface = gr.Interface(
    fn=chatbot_response,
    inputs=[pdf_input, question_input],
    outputs=response_output,
    title="RAG Chatbot mit PDF-Unterstützung",
    description="Lade eine PDF-Datei hoch und stelle Fragen zu ihrem Inhalt."
)

if __name__ == "__main__":
    interface.launch()