# app.py
import gradio as gr
from transformers import pipeline, T5Tokenizer, AutoModelForSeq2SeqLM

# ====== Summarization Model ======
summarizer = pipeline("summarization", model="facebook/bart-large-cnn")

# ====== Preload T5 Model & Tokenizer for Quiz / Flashcards ======
tokenizer = T5Tokenizer.from_pretrained("iarfmoose/t5-base-question-generator")
model = AutoModelForSeq2SeqLM.from_pretrained("iarfmoose/t5-base-question-generator")
question_generator = pipeline("text2text-generation", model=model, tokenizer=tokenizer)

# ====== Summarizer Function ======
def summarize_text(text, mode):
    if not text.strip():
        return "⚠️ Please provide some notes!"
    
    if mode == "short":
        max_len, min_len = 60, 20
    elif mode == "revision":
        max_len, min_len = 120, 40
    else:  # long
        max_len, min_len = 250, 80

    try:
        summary = summarizer(
            text, 
            max_length=max_len, 
            min_length=min_len, 
            do_sample=False,
            truncation=True
        )
        return summary[0]['summary_text']
    except Exception as e:
        return f"❌ Error: {str(e)}"

# ====== Quiz Generator ======
def generate_quiz(text, num_questions):
    if not text.strip():
        return "⚠️ Please provide some notes!"
    
    prompt = f"Generate {num_questions} quiz questions from the following text:\n{text}"
    output = question_generator(prompt, max_length=512)
    return output[0]['generated_text']

# ====== Flashcards Generator ======
def generate_flashcards(text, num_flashcards):
    if not text.strip():
        return "⚠️ Please provide some notes!"
    
    prompt = f"Generate {num_flashcards} flashcards (Q&A) from the following text:\n{text}"
    output = question_generator(prompt, max_length=512)
    return output[0]['generated_text']

# ====== Gradio Interface ======
with gr.Blocks() as demo:
    gr.Markdown("## 📝 AI Notes Summarizer, Quiz & Flashcards")
    
    # --- Summarizer Tab ---
    with gr.Tab("Summarizer"):
        input_text = gr.Textbox(lines=10, placeholder="Paste your notes here...")
        mode = gr.Radio(["short", "revision", "long"], value="revision", label="Summary Type")
        summary_output = gr.Textbox(label="Summary")
        gr.Button("Summarize").click(fn=summarize_text, inputs=[input_text, mode], outputs=[summary_output])
    
    # --- Quiz Generator Tab ---
    with gr.Tab("Quiz Generator"):
        quiz_input = gr.Textbox(lines=10, placeholder="Paste your notes or summary here...")
        num_q = gr.Dropdown(["5","10","15","20","25","30"], value="5", label="Number of Questions")
        quiz_output = gr.Textbox(label="Quiz Questions")
        gr.Button("Generate Quiz").click(fn=generate_quiz, inputs=[quiz_input, num_q], outputs=[quiz_output])
    
    # --- Flashcards Generator Tab ---
    with gr.Tab("Flashcards Generator"):
        flash_input = gr.Textbox(lines=10, placeholder="Paste your notes or summary here...")
        num_f = gr.Dropdown(["5","10","15","20","25","30"], value="5", label="Number of Flashcards")
        flash_output = gr.Textbox(label="Flashcards (Q&A)")
        gr.Button("Generate Flashcards").click(fn=generate_flashcards, inputs=[flash_input, num_f], outputs=[flash_output])

demo.launch()