import torch
from transformers import RagRetriever, RagTokenizer, RagSequenceForGeneration
import gradio as gr
from PyPDF2 import PdfReader
import re

# Laden des Tokenizers, des Retrievers und des Modells (auf CPU)
model_name = "facebook/rag-token-nq"
tokenizer = RagTokenizer.from_pretrained(model_name)
retriever = RagRetriever.from_pretrained(model_name, index_name="exact")
model = RagSequenceForGeneration.from_pretrained(model_name).to("cpu")  # Modell auf CPU laden

# Funktion zum Extrahieren und Bereinigen von Text aus PDF
def extract_text_from_pdf(pdf_path):
    reader = PdfReader(pdf_path)
    text = ""
    for page in reader.pages:
        page_text = page.extract_text()
        if page_text:
            text += page_text
    return text

def clean_text(text):
    # Entfernen unnötiger Zeichen, Reduktion von Leerzeichen
    text = re.sub(r'\s+', ' ', text)
    text = re.sub(r'[^\w\s.,-]', '', text)
    return text.strip()

# Funktion zum Aufteilen langer Texte in Abschnitte
def split_text_into_chunks(text, chunk_size=1000):
    words = text.split()
    chunks = [' '.join(words[i:i + chunk_size]) for i in range(0, len(words), chunk_size)]
    return chunks

# Hauptfunktion für die Fragebeantwortung mit RAG
def chatbot_response(pdf_path, question):
    try:
        # PDF-Inhalt extrahieren und bereinigen
        context = clean_text(extract_text_from_pdf(pdf_path))
        if not context:
            return "Das Dokument enthält keinen Text oder konnte nicht gelesen werden."

        # Dokumenttext in Abschnitte aufteilen, um Speicher zu sparen
        chunks = split_text_into_chunks(context)

        # Antwortgenerierung mit minimalem Speicherverbrauch
        answers = []
        with torch.no_grad():  # Verhindert das Speichern von Gradienten (für CPU wichtig)
            for chunk in chunks:
                retriever.index = [chunk]
                inputs = tokenizer(question, return_tensors="pt").to("cpu")  # Sicherstellen, dass Inputs auf CPU bleiben
                generated_ids = model.generate(**inputs, max_length=150)  # Kürzere Antwortlänge
                answer = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
                if answer:
                    answers.append(answer)

        final_answer = " / ".join(answers) if answers else "Keine spezifische Antwort gefunden."
        return final_answer

    except Exception as e:
        return f"Es ist ein Fehler aufgetreten: {str(e)}"

# Gradio-Interface erstellen
pdf_input = gr.File(label="PDF-Datei hochladen", type="filepath")
question_input = gr.Textbox(label="Frage eingeben", placeholder="Stelle eine Frage zu dem PDF-Dokument")
response_output = gr.Textbox(label="Antwort", lines=4)

interface = gr.Interface(
    fn=chatbot_response,
    inputs=[pdf_input, question_input],
    outputs=response_output,
    title="RAG PDF-Fragebeantwortung auf CPU",
    description="Lade eine PDF-Datei hoch und stelle Fragen zu ihrem Inhalt. Das System verwendet Retrieval-Augmented Generation (RAG) auf CPU zur Beantwortung.",
)

# Interface für Hugging Face Spaces
interface.launch(share=True)