import gradio as gr
from transformers import pipeline

gr.Markdown("## 🎯 Smart Learning Assistant")
# Load models (lightweight for Spaces)
qa_pipeline = pipeline("text2text-generation", model="google/flan-t5-small")
summarizer = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6")

# ---------------- FUNCTIONS ---------------- #

def answer_question(text):
    result = qa_pipeline(f"Answer this question: {text}", max_length=100)
    return result[0]['generated_text']


def summarize_text(text):
    result = summarizer(text, max_length=80, min_length=20, do_sample=False)
    return result[0]['summary_text']


def generate_quiz(topic):
    prompt = f"Generate 3 multiple choice questions with answers on: {topic}"
    result = qa_pipeline(prompt, max_length=200)
    return result[0]['generated_text']


# ---------------- UI ---------------- #

with gr.Blocks(theme=gr.themes.Soft()) as demo:

    gr.Markdown("# 📚 AI Study Assistant")
    gr.Markdown("### 🚀 Ask questions | Summarize notes | Generate quizzes")

    with gr.Tabs():

        # 🔹 TAB 1: Q&A
        with gr.Tab("💬 Ask Doubt"):
            q_input = gr.Textbox(lines=3, placeholder="Enter your question...")
            q_output = gr.Textbox(label="Answer")
            q_btn = gr.Button("Get Answer")
            q_btn.click(answer_question, inputs=q_input, outputs=q_output)

        # 🔹 TAB 2: Summary
        with gr.Tab("📝 Summarize"):
            s_input = gr.Textbox(lines=6, placeholder="Paste your notes here...")
            s_output = gr.Textbox(label="Summary")
            s_btn = gr.Button("Summarize")
            s_btn.click(summarize_text, inputs=s_input, outputs=s_output)

        # 🔹 TAB 3: Quiz
        with gr.Tab("🧠 Generate Quiz"):
            quiz_input = gr.Textbox(placeholder="Enter topic (e.g., DBMS, AI)")
            quiz_output = gr.Textbox(label="Quiz")
            quiz_btn = gr.Button("Generate")
            quiz_btn.click(generate_quiz, inputs=quiz_input, outputs=quiz_output)

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