import gradio as gr
from transformers import pipeline

# ========================
# Load models (tối ưu cho HF Free + đa ngôn ngữ)
# ========================

summarizer = pipeline(
    "summarization",
    model="csebuetnlp/mT5_multilingual_XLSum"
)

quiz_generator = pipeline(
    "text2text-generation",
    model="google/mt5-small"
)

# ========================
# AI Functions
# ========================

def summarize_and_quiz(text, num_questions):

    num_questions = int(num_questions)

    if not text or len(text) < 100:
        return "Văn bản quá ngắn. Hãy nhập ít nhất 100 ký tự.", ""

    # Cắt bớt để tránh quá giới hạn model
    text = text[:3000]

    # --- TÓM TẮT ---
    try:
        summary = summarizer(
            text,
            max_length=200,
            min_length=80,
            do_sample=False
        )[0]['summary_text']
    except:
        summary = "❌ Có lỗi xảy ra trong quá trình tóm tắt."

    # --- TẠO QUIZ ---
    prompt = f"""
Dựa trên nội dung sau, hãy tạo {num_questions} câu hỏi trắc nghiệm.

NỘI DUNG:
{summary}

YÊU CẦU:
1. Mỗi câu hỏi nằm trên 1 dòng riêng
2. Có 4 đáp án A, B, C, D
3. Phải ghi đáp án đúng ở cuối câu
4. Không sử dụng ký hiệu như <extra_id_0> hoặc <extra_id_1>
5. Viết cùng ngôn ngữ với nội dung (Việt/Anh)

Ví dụ:
1. Câu hỏi ...?
A. ...
B. ...
C. ...
D. ...
Đáp án: B

BẮT ĐẦU TẠO CÂU HỎI:
"""

    try:
        quiz = quiz_generator(
            prompt,
            max_length=512,
            do_sample=False
        )[0]["generated_text"]

        # Xoá hoàn toàn rác của T5
        quiz = quiz.replace("<extra_id_0>", "").replace("<extra_id_1>", "").strip()

        if len(quiz) < 30:
            quiz = "❌ Không tạo được câu hỏi phù hợp. Hãy thử lại với nội dung dài hơn."

    except:
        quiz = "❌ Có lỗi xảy ra trong quá trình tạo câu hỏi."

    return summary, quiz


# ========================
# Gradio Interface
# ========================

interface = gr.Interface(
    fn=summarize_and_quiz,

    inputs=[
        gr.Textbox(
            label="📄 Nhập văn bản (Tiếng Việt hoặc Tiếng Anh)",
            placeholder="Dán nội dung bài học / tài liệu / văn bản cần tóm tắt vào đây...",
            lines=15
        ),
        gr.Slider(5, 10, value=5, step=1, label="Số câu hỏi (Quiz)")
    ],

    outputs=[
        gr.Textbox(label="📌 Nội dung tóm tắt", lines=8),
        gr.Textbox(label="📝 Câu hỏi ôn tập (Quiz)", lines=12)
    ],

    title="📚 AI TÓM TẮT & TẠO QUIZ CHO SINH VIÊN",
    description="""
    ✅ Hỗ trợ: Tiếng Việt + Tiếng Anh  
    ✅ Đã fix lỗi: <extra_id_0>  
    ✅ Tóm tắt không mất ý chính  
    ✅ Tạo 5–10 câu hỏi trắc nghiệm tự động  
    ✅ Tối ưu chạy trên HuggingFace Free  

    👉 Chỉ cần dán nội dung vào và nhấn Submit.
    """
)

interface.launch()