import gradio as gr
import spacy
import torch
import subprocess
import sys
import tempfile
import os
from transformers import T5ForConditionalGeneration, AutoTokenizer
from fpdf import FPDF

subprocess.run(
    [sys.executable, "-m", "spacy", "download", "en_core_web_sm"],
    check=True
)

MODEL_ID = "Hamzasajjad38/t5-small-qg"
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=True)
model     = T5ForConditionalGeneration.from_pretrained(MODEL_ID)
device    = "cuda" if torch.cuda.is_available() else "cpu"
model     = model.to(device)

nlp = spacy.load("en_core_web_sm")

# ── Core generation function ──────────────────────────────────────────
def generate_questions(passage, num_questions, question_type):
    if not passage.strip():
        return "⚠️ Please enter a passage.", None

    doc = nlp(passage)

    # Extract candidate answers: named entities first, then noun chunks
    candidates = [ent.text for ent in doc.ents]
    if len(candidates) < num_questions:
        candidates += [chunk.text for chunk in doc.noun_chunks]

    # Deduplicate and limit
    seen = set()
    unique = []
    for c in candidates:
        if c.lower() not in seen and len(c.split()) <= 5:
            seen.add(c.lower())
            unique.append(c)
    candidates = unique[:num_questions]

    if not candidates:
        return "⚠️ Could not extract answer candidates from this passage. Try a longer or more detailed passage.", None

    output_lines = []
    for i, answer in enumerate(candidates):
        highlighted = passage.replace(answer, f"<hl> {answer} <hl>", 1)
        input_text  = f"generate question: {highlighted}"
        inputs  = tokenizer(input_text, return_tensors="pt",
                            max_length=512, truncation=True).to(device)
        outputs = model.generate(**inputs, max_new_tokens=64,
                                 num_beams=4, early_stopping=True)
        question = tokenizer.decode(outputs[0], skip_special_tokens=True)

        if question_type == "True / False":
            question = f"True or False: {question.rstrip('?')}?"

        output_lines.append(f"Q{i+1}: {question}")
        output_lines.append(f"   Answer: {answer}\n")

    return "\n".join(output_lines)

# ── PDF export ────────────────────────────────────────────────────────
def export_pdf(passage, num_questions, question_type):
    result = generate_questions(passage, num_questions, question_type)
    if result.startswith("⚠️"):
        return None

    pdf = FPDF()
    pdf.add_page()
    pdf.set_font("Helvetica", "B", 14)
    pdf.cell(0, 10, "Generated Question Paper", ln=True, align="C")
    pdf.ln(4)
    pdf.set_font("Helvetica", size=10)
    pdf.multi_cell(0, 6, f"Passage: {passage[:200]}...")
    pdf.ln(6)
    pdf.set_font("Helvetica", size=11)
    for line in result.split("\n"):
        pdf.multi_cell(0, 7, line)
    pdf.ln(4)

    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".pdf")
    pdf.output(tmp.name)
    return tmp.name

# ── Gradio UI ─────────────────────────────────────────────────────────
with gr.Blocks(title="Automatic Question Generator", theme=gr.themes.Soft()) as demo:

    gr.Markdown("""
    # 🧠 Automatic Question Generator
    **Fine-tuned T5-small on SQuAD 1.1** · Generates questions from any educational passage
    """)

    with gr.Row():
        with gr.Column(scale=2):
            passage_input = gr.Textbox(
                label="Input Passage",
                placeholder="Paste any educational paragraph here...",
                lines=8
            )
            with gr.Row():
                num_q = gr.Slider(1, 6, value=3, step=1, label="Number of questions")
                q_type = gr.Dropdown(
                    ["Short answer", "True / False"],
                    value="Short answer",
                    label="Question type"
                )
            with gr.Row():
                gen_btn = gr.Button("⚡ Generate Questions", variant="primary")
                pdf_btn = gr.Button("📄 Export PDF")

        with gr.Column(scale=2):
            output_box = gr.Textbox(label="Generated Questions", lines=16, interactive=False)
            pdf_output = gr.File(label="Download PDF", visible=True)

    # Sample passages
    gr.Examples(
        examples=[
            ["Photosynthesis is a process used by plants to convert light energy into chemical energy stored in glucose. It occurs inside chloroplasts using chlorophyll to absorb sunlight.", 3, "Short answer"],
            ["The Industrial Revolution began in Britain in the late 18th century. The steam engine, invented by James Watt, transformed manufacturing, agriculture, and transportation across the country.", 3, "Short answer"],
            ["Machine learning is a subset of artificial intelligence where systems learn from data to improve performance without being explicitly programmed.", 2, "True / False"],
        ],
        inputs=[passage_input, num_q, q_type]
    )

    gen_btn.click(
        fn=generate_questions,
        inputs=[passage_input, num_q, q_type],
        outputs=output_box
    )
    pdf_btn.click(
        fn=export_pdf,
        inputs=[passage_input, num_q, q_type],
        outputs=pdf_output
    )

demo.launch(share=False, server_name="0.0.0.0")