import gradio as gr
from groq import Groq
from docx import Document
import datetime
import os

# ✅ Initialize Groq Client
client = Groq(api_key="YOUR_GROQ_API_KEY")

# ================== Helper Functions ==================

def generate_text(prompt, model_name):
    response = client.chat.completions.create(
        messages=[{"role": "user", "content": prompt}],
        model=model_name,
    )
    return response.choices[0].message.content.strip()

def save_docx(content, filename):
    doc = Document()
    doc.add_paragraph(content)
    filepath = f"/tmp/{filename}.docx"
    doc.save(filepath)
    return filepath

def save_md(content, filename):
    content = content.replace("✨", "").replace("🌟", "").replace("✅", "")
    filepath = f"/tmp/{filename}.md"
    with open(filepath, "w", encoding="utf-8") as f:
        f.write(content)
    return filepath

# ================== Core Generation Functions ==================

def generate_lesson_plan(language, grade, subject, num_students, date, slo, content):
    if language == "English":
        model = "llama-3.1-8b-instant"
        prompt = f"""
        Create a SMART Lesson Plan in English for:
        Grade: {grade}
        Subject: {subject}
        Number of Students: {num_students}
        Date: {date}

        Include the following sections:
        1. Input Table
        2. SLOs
        3. Content to be Delivered
        4. Home Assignment (aligned with SLOs)
        5. Feedback by Mentor/Principal

        Use this context:
        - SLOs: {slo}
        - Content: {content}
        """
    else:
        model = "openai/gpt-oss-120b"
        prompt = f"""
        درج ذیل معلومات کی بنیاد پر ایک SMART سبق کا منصوبہ اردو میں تیار کریں۔
        جماعت: {grade}
        مضمون: {subject}
        طلبہ کی تعداد: {num_students}
        تاریخ: {date}

        درج ذیل حصے شامل ہوں:
        1. ان پٹ جدول  
        2. سیکھنے کے مقاصد (SLOs)  
        3. پڑھائی جانے والا مواد  
        4. گھریلو مشق (جو SLOs سے ہم آہنگ ہو)  
        5. سرپرست / پرنسپل کی رائے  

        سیاق و سباق:
        - سیکھنے کے مقاصد: {slo}
        - مواد: {content}
        """

    output = generate_text(prompt, model)
    filename = f"SMART_Lesson_Plan_{language}_{grade}_{subject}"
    if language == "English":
        filepath = save_docx(output, filename)
    else:
        filepath = save_md(output, filename)
    return output, filepath


def generate_flashcards(language, grade, subject, slo, content):
    if language == "English":
        model = "llama-3.1-8b-instant"
        prompt = f"""
        Create concise classroom flashcards in English for Grade {grade}, Subject {subject}.
        Align flashcards with the following SLOs and lesson content:
        SLOs: {slo}
        Content: {content}
        Format as Q&A flashcards.
        """
    else:
        model = "openai/gpt-oss-120b"
        prompt = f"""
        درج ذیل سیکھنے کے مقاصد اور مواد کی بنیاد پر جماعت {grade} کے لیے فلیش کارڈز تیار کریں۔
        مضمون: {subject}
        ہر کارڈ میں سوال اور مختصر جواب شامل ہو۔
        سیکھنے کے مقاصد: {slo}
        مواد: {content}
        """

    output = generate_text(prompt, model)
    filename = f"Flashcards_{language}_{grade}_{subject}"
    if language == "English":
        filepath = save_docx(output, filename)
    else:
        filepath = save_md(output, filename)
    return output, filepath


def generate_quiz(language, grade, subject, slo, content):
    if language == "English":
        model = "llama-3.1-8b-instant"
        prompt = f"""
        Create a short quiz in English for Grade {grade}, Subject {subject}.
        Align questions with the following SLOs and lesson content.
        SLOs: {slo}
        Content: {content}
        Include 5 multiple-choice questions with answers.
        """
    else:
        model = "openai/gpt-oss-120b"
        prompt = f"""
        جماعت {grade}، مضمون {subject} کے لیے ایک مختصر کوئز اردو میں تیار کریں۔
        کوئز کے سوالات درج ذیل سیکھنے کے مقاصد اور مواد سے ہم آہنگ ہوں۔
        سیکھنے کے مقاصد: {slo}
        مواد: {content}
        پانچ کثیرالاختیاری سوالات (MCQs) شامل کریں اور جوابات بھی دیں۔
        """

    output = generate_text(prompt, model)
    filename = f"Quiz_{language}_{grade}_{subject}"
    if language == "English":
        filepath = save_docx(output, filename)
    else:
        filepath = save_md(output, filename)
    return output, filepath


# ================== Gradio App ==================

def main_app(language, feature, grade, subject, num_students, date, slo, content):
    if feature == "Lesson Plan":
        result, file = generate_lesson_plan(language, grade, subject, num_students, date, slo, content)
        return result, file, None, None

    elif feature == "Flashcards":
        result, file = generate_flashcards(language, grade, subject, slo, content)
        return result, None, file, None

    elif feature == "Quiz":
        result, file = generate_quiz(language, grade, subject, slo, content)
        return result, None, None, file

    else:
        return "Please select a valid feature.", None, None, None


# ================== UI Layout ==================

with gr.Blocks(title="SMART Lesson Plan Generator") as app:
    gr.Markdown("<h1 style='text-align:center;'>📘 SMART LESSON PLAN GENERATOR</h1>")

    with gr.Row():
        language = gr.Dropdown(["English", "Urdu"], label="Select Language", value="English")
        feature = gr.Dropdown(["Lesson Plan", "Flashcards", "Quiz"], label="Select Feature", value="Lesson Plan")

    with gr.Row():
        grade = gr.Dropdown(
            ["Grade 1", "Grade 2", "Grade 3", "Grade 4", "Grade 5", "Grade 6", 
             "Grade 7", "Grade 8", "Grade 9", "Grade 10", "Grade 11", "Grade 12"],
            label="Select Class/Grade"
        )
        subject = gr.Dropdown(
            ["English", "Urdu", "Mathematics", "Science", "Islamiyat", "Social Studies", "Computer Science"],
            label="Select Subject"
        )

    with gr.Row():
        num_students = gr.Number(label="Number of Students", value=30)
        date = gr.Textbox(label="Date (YYYY-MM-DD)", value=str(datetime.date.today()))

    slo = gr.Textbox(label="Enter SLOs (Learning Outcomes)", lines=3)
    content = gr.Textbox(label="Enter Lesson Content Summary", lines=3)

    generate_btn = gr.Button("🚀 Generate")
    output_text = gr.Textbox(label="Generated Output", lines=20)

    lesson_plan_file = gr.File(label="📄 Download Lesson Plan")
    flashcard_file = gr.File(label="📘 Download Flashcards")
    quiz_file = gr.File(label="🧾 Download Quiz")

    generate_btn.click(
        main_app,
        inputs=[language, feature, grade, subject, num_students, date, slo, content],
        outputs=[output_text, lesson_plan_file, flashcard_file, quiz_file]
    )

app.launch()