import gradio as gr
import pdfplumber
import torch
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# ------------------------------------------------------------------------------
# 1) Modell laden (valhalla/t5-base-qg-hl)
# ------------------------------------------------------------------------------
MODEL_NAME = "valhalla/t5-base-qg-hl"
print(f"Lade Tokenizer und Modell: {MODEL_NAME}")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)

# Optional: auf CPU behalten (gratis, aber langsamer) oder GPU nutzen, wenn Hugging Face
# Community GPU verfügbar ist (dann model.to("cuda"))
device = torch.device("cpu")
model.to(device)


# ------------------------------------------------------------------------------
# 2) PDF-Text extrahieren
# ------------------------------------------------------------------------------
def extract_text_from_pdf(pdf_file):
    """
    Extrahiert den Text aller Seiten aus einem hochgeladenen PDF.
    """
    if pdf_file is None:
        return ""

    text = ""
    with pdfplumber.open(pdf_file) as pdf:
        for page in pdf.pages:
            text += page.extract_text() + "\n"

    # Minimales Cleaning
    text = " ".join(text.split())
    return text

# ------------------------------------------------------------------------------
# 3) Question Generation mit T5
#    Dieses Modell ("valhalla/t5-base-qg-hl") nutzt einen 'Highlight-basierten'
#    Ansatz. Am einfachsten probieren wir, den gesamten Text an das Modell zu geben.
#    Für bessere Qualität könnte man (a) Text kürzen, (b) "answer highlighting" machen.
# ------------------------------------------------------------------------------
def generate_questions(text_chunk, max_length=128):
    """
    Fragt das T5-QG-Modell nach Fragen für den gegebenen Text.
    Achtung: 'valhalla/t5-base-qg-hl' erwartet i.d.R. 'question: ... context: ...' oder
    'generate question: ...' Prompts. Wir machen ein einfaches prompt-engineering.
    """
    # Einfacher Prompt: wir fügen "generate question:" voran
    prompt_text = f"generate question: {text_chunk}"

    inputs = tokenizer.encode(prompt_text, return_tensors="pt").to(device)
    output = model.generate(
        inputs, 
        max_length=max_length, 
        num_beams=4, 
        early_stopping=True
    )
    question = tokenizer.decode(output[0], skip_special_tokens=True)
    return question


# ------------------------------------------------------------------------------
# 4) Gradio-Workflows
# ------------------------------------------------------------------------------
def process_pdf(pdf_file, num_questions):
    """
    1) PDF extrahieren
    2) Kürzen (Text chunk) 
    3) Mehrere Fragen generieren
    """
    if pdf_file is None:
        return "Keine PDF hochgeladen."

    # PDF-Text holen
    text = extract_text_from_pdf(pdf_file.name)
    if not text:
        return "Text konnte nicht extrahiert werden oder PDF ist leer."

    # Ggf. nur ersten 1500 Zeichen nehmen, damit wir das Modell nicht überfüttern
    text_chunk = text[:1500]

    # Generiere mehrere Fragen
    questions_output = []
    for i in range(num_questions):
        q = generate_questions(text_chunk)
        questions_output.append(f"Frage {i+1}: {q}")

    # Kombiniere das als Ausgabe
    final_output = "\n\n".join(questions_output)
    return final_output

def build_app():
    with gr.Blocks() as demo:
        gr.Markdown("# QG-PDF – Fragegenerierung aus PDF (ohne OpenAI)")
        gr.Markdown(
            "Lade ein PDF hoch und wähle, wie viele Fragen generiert werden sollen. "
            "Wir verwenden das Modell **valhalla/t5-base-qg-hl**, "
            "das (meist) eine offene Frage ausgibt."
        )

        with gr.Row():
            pdf_file = gr.File(label="📄 PDF hochladen")
            q_slider = gr.Slider(
                minimum=1,
                maximum=5,
                step=1,
                value=3,
                label="Anzahl Fragen"
            )
        generate_btn = gr.Button("Fragen generieren")
        output_box = gr.Textbox(
            label="Generierte Fragen",
            lines=10
        )

        def on_click_generate(pdf, q_num):
            return process_pdf(pdf, q_num)

        generate_btn.click(
            fn=on_click_generate,
            inputs=[pdf_file, q_slider],
            outputs=[output_box]
        )

    return demo

demo = build_app()

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