import gradio as gr
import easyocr
from pdf2image import convert_from_path
from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration
import os

# Initialisiere EasyOCR für Deutsch
reader = easyocr.Reader(['de'])  # für die deutsche Sprache

# Initialisiere das deutsche Modell und den Tokenizer für RAG
model_name = "deepset/gbert-base"  # Beispiel für ein deutsches Modell
tokenizer = RagTokenizer.from_pretrained(model_name)
model = RagSequenceForGeneration.from_pretrained(model_name)
retriever = RagRetriever.from_pretrained(model_name, index_name="exact", use_dummy_dataset=True)

# OCR-Funktion: Konvertiert PDF zu Bildern und extrahiert Text mit EasyOCR
def extract_text_from_pdf(file):
    # Konvertiere PDF-Seiten in Bilder
    images = convert_from_path(file.name, 300)  # 300 DPI für bessere Qualität
    
    text = ""
    # Extrahiere Text aus jedem Bild mit EasyOCR
    for image in images:
        ocr_result = reader.readtext(image)
        for detection in ocr_result:
            text += detection[1] + "\n"
    
    return text

# Funktion zur Generierung einer Antwort basierend auf dem hochgeladenen Dokument
def get_rag_answer(input_message, uploaded_file):
    # Extrahiere den Text aus dem hochgeladenen PDF-Dokument mit OCR
    document_text = extract_text_from_pdf(uploaded_file)

    # Simuliere den Retrieval-Mechanismus, indem wir den extrahierten Text verwenden
    inputs = tokenizer(input_message, return_tensors="pt")
    retrieved_docs = retriever.retrieve(input_ids=inputs["input_ids"])

    # Kombiniere die extrahierten Dokumente und frage das Modell zur Generierung einer Antwort
    input_ids = tokenizer(input_message, return_tensors="pt").input_ids
    generated_ids = model.generate(input_ids=input_ids, 
                                   decoder_start_token_id=model.config.pad_token_id, 
                                   num_beams=4, 
                                   max_length=100)

    answer = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
    
    # Hier könnten wir eine Referenz (z.B. Absatz, Seite) in die Antwort einfügen
    references = "Referenz: Abschnitt X, Seite Y (aus Dokument)"  # Füge diese Infos hinzu, falls möglich

    return f"{answer} \n\n{references}"

# Gradio-Oberfläche
def gradio_interface():
    iface = gr.Interface(
        fn=get_rag_answer,
        inputs=[
            gr.Textbox(label="User Input", placeholder="Stellen Sie eine Frage..."),
            gr.File(label="Laden Sie ein PDF-Dokument hoch", type="file")  # Ermöglicht das Hochladen von PDF-Dateien
        ],
        outputs=gr.Textbox(label="Antwort des Chatbots"),
        live=True  # Sofortige Antwortgenerierung
    )
    iface.launch()

# Starte die Gradio-Oberfläche
gradio_interface()