import os
import gradio as gr
from docx import Document
from docx.shared import Pt, Inches, RGBColor
from docx.enum.text import WD_ALIGN_PARAGRAPH
import datetime
from groq import Groq

# Get Groq API key from environment variable
GROQ_API_KEY = os.getenv("GROQ_API_KEY")
client = Groq(api_key=GROQ_API_KEY)

# Counter for lesson plans generated
lesson_plan_counter = 0

# Pakistan National Curriculum 2006 - Subject Lists by Grade
SUBJECTS_BY_GRADE = {
    "Grade 1": ["English", "Urdu", "Mathematics", "General Science", "Islamiyat", "Drawing"],
    "Grade 2": ["English", "Urdu", "Mathematics", "General Science", "Islamiyat", "Drawing"],
    "Grade 3": ["English", "Urdu", "Mathematics", "General Science", "Social Studies", "Islamiyat", "Drawing"],
    "Grade 4": ["English", "Urdu", "Mathematics", "General Science", "Social Studies", "Islamiyat", "Drawing"],
    "Grade 5": ["English", "Urdu", "Mathematics", "General Science", "Social Studies", "Islamiyat", "Drawing"],
    "Grade 6": ["English", "Urdu", "Mathematics", "Science", "Social Studies", "Islamiyat", "Computer Science"],
    "Grade 7": ["English", "Urdu", "Mathematics", "Science", "Social Studies", "Islamiyat", "Computer Science"],
    "Grade 8": ["English", "Urdu", "Mathematics", "Science", "Social Studies", "Islamiyat", "Computer Science"],
    "Grade 9": ["English", "Urdu", "Mathematics", "Physics", "Chemistry", "Biology", "Pakistan Studies", "Islamiyat", "Computer Science"],
    "Grade 10": ["English", "Urdu", "Mathematics", "Physics", "Chemistry", "Biology", "Pakistan Studies", "Islamiyat", "Computer Science"],
    "Grade 11": {
        "Pre-Medical": ["English", "Urdu", "Physics", "Chemistry", "Biology", "Islamiyat", "Pakistan Studies"],
        "Pre-Engineering": ["English", "Urdu", "Physics", "Chemistry", "Mathematics", "Islamiyat", "Pakistan Studies", "Computer Science"],
        "Computer Science": ["English", "Urdu", "Physics", "Mathematics", "Computer Science", "Islamiyat", "Pakistan Studies"],
        "Humanities": ["English", "Urdu", "Pakistan Studies", "Islamiyat", "Economics", "Psychology", "Education", "Sociology"]
    },
    "Grade 12": {
        "Pre-Medical": ["English", "Urdu", "Physics", "Chemistry", "Biology", "Islamiyat", "Pakistan Studies"],
        "Pre-Engineering": ["English", "Urdu", "Physics", "Chemistry", "Mathematics", "Islamiyat", "Pakistan Studies", "Computer Science"],
        "Computer Science": ["English", "Urdu", "Physics", "Mathematics", "Computer Science", "Islamiyat", "Pakistan Studies"],
        "Humanities": ["English", "Urdu", "Pakistan Studies", "Islamiyat", "Economics", "Psychology", "Education", "Sociology"]
    }
}

def get_subjects_for_grade(grade):
    """Return list of subjects based on selected grade"""
    if grade in ["Grade 11", "Grade 12"]:
        all_subjects = set()
        for stream_subjects in SUBJECTS_BY_GRADE[grade].values():
            all_subjects.update(stream_subjects)
        return sorted(list(all_subjects))
    return SUBJECTS_BY_GRADE.get(grade, [])

def generate_lesson_plan(school_name, grade, subject, topic, students, duration, date, teacher):
    global lesson_plan_counter
    
    prompt = f"""
You are a professional lesson plan expert creating detailed instructional materials.
Create a comprehensive English lesson plan using the BOPPPS model for the topic: "{topic}" for {grade} {subject}.
Do not use emojis, asterisks, or special characters in your response.
Do not repeat the basic information in headings or body.
Provide ONLY the following structured components with these EXACT headings:

Bridge-In:
Write a short, engaging introductory paragraph that connects students' prior knowledge to the new topic.

Learning Objectives:
List exactly 3 clear, measurable learning objectives using action verbs.

Pre-Assessment Questions:
Provide exactly 2 diagnostic questions to gauge students' prior understanding.

Participatory Learning Activity:
Describe one detailed, interactive learning activity with clear step-by-step instructions.

Post-Assessment Questions:
Provide 2-3 evaluation questions to assess understanding of the lesson objectives.

Summary:
Write a concise paragraph summarizing the key concepts covered in the lesson.

Home Assignment:
Provide a meaningful assignment related to the topic that students can complete and present in the next class. Include clear instructions and expected outcomes.

Format your response professionally without using any special formatting characters. Use the exact heading names provided above.
"""
    
    try:
        response = client.chat.completions.create(
            model="llama-3.1-8b-instant",
            messages=[{"role": "user", "content": prompt}],
            temperature=0.7,
        )
        lesson_text = response.choices[0].message.content.strip()
        
        # Create Word document with professional formatting
        doc = Document()
        
        # Set document margins
        sections = doc.sections
        for section in sections:
            section.top_margin = Inches(0.75)
            section.bottom_margin = Inches(0.75)
            section.left_margin = Inches(1)
            section.right_margin = Inches(1)
        
        # Add school name as header (centered, uppercase)
        header = doc.sections[0].header
        header_para = header.paragraphs[0]
        header_para.text = school_name.upper()
        header_para.alignment = WD_ALIGN_PARAGRAPH.CENTER
        header_run = header_para.runs[0]
        header_run.font.size = Pt(14)
        header_run.font.bold = True
        header_run.font.color.rgb = RGBColor(0, 0, 0)
        
        # Add page numbers to footer (just numbers, no "Page" text)
        footer = doc.sections[0].footer
        footer_para = footer.paragraphs[0]
        footer_para.alignment = WD_ALIGN_PARAGRAPH.CENTER
        # Add field code for page number
        from docx.oxml import OxmlElement
        from docx.oxml.ns import qn
        
        run = footer_para.add_run()
        fldChar1 = OxmlElement('w:fldChar')
        fldChar1.set(qn('w:fldCharType'), 'begin')
        
        instrText = OxmlElement('w:instrText')
        instrText.set(qn('xml:space'), 'preserve')
        instrText.text = 'PAGE'
        
        fldChar2 = OxmlElement('w:fldChar')
        fldChar2.set(qn('w:fldCharType'), 'end')
        
        run._r.append(fldChar1)
        run._r.append(instrText)
        run._r.append(fldChar2)
        
        # Metadata table
        table = doc.add_table(rows=0, cols=2)
        table.style = 'Light Grid Accent 1'
        
        data = {
            "Grade": grade,
            "Subject": subject,
            "Topic": topic,
            "Number of Students": students,
            "Duration": f"{duration} minutes",
            "Date": date,
            "Teacher Name": teacher
        }
        
        for key, value in data.items():
            row_cells = table.add_row().cells
            row_cells[0].text = key
            row_cells[1].text = str(value)
            row_cells[0].paragraphs[0].runs[0].font.bold = True
        
        # Add spacing
        doc.add_paragraph()
        
        # Process lesson text to make headings bold
        headings_to_bold = [
            "Bridge-In:",
            "Learning Objectives:",
            "Pre-Assessment Questions:",
            "Participatory Learning Activity:",
            "Post-Assessment Questions:",
            "Summary:",
            "Home Assignment:"
        ]
        
        # Split content by lines and process
        lines = lesson_text.split('\n')
        for line in lines:
            line = line.strip()
            if not line:
                doc.add_paragraph()
                continue
            
            # Check if line starts with any heading
            is_heading = False
            for heading in headings_to_bold:
                if line.startswith(heading):
                    para = doc.add_paragraph()
                    run = para.add_run(line)
                    run.font.bold = True
                    is_heading = True
                    break
            
            if not is_heading:
                doc.add_paragraph(line)
        
        # Save document
        filename = f"lesson_plan_{datetime.datetime.now().strftime('%Y%m%d%H%M%S')}.docx"
        filepath = f"/tmp/{filename}"
        doc.save(filepath)
        
        # Increment counter
        lesson_plan_counter += 1
        
        return filepath, lesson_plan_counter
        
    except Exception as e:
        # Return an error document
        error_doc = Document()
        error_doc.add_heading("Error in Lesson Plan Generator", 0)
        error_doc.add_paragraph(f"An error occurred: {str(e)}")
        error_doc.add_paragraph("Please check your inputs and try again.")
        filename = f"error_{datetime.datetime.now().strftime('%Y%m%d%H%M%S')}.docx"
        filepath = f"/tmp/{filename}"
        error_doc.save(filepath)
        return filepath, lesson_plan_counter

def update_subjects(grade):
    """Update subject dropdown based on selected grade"""
    subjects = get_subjects_for_grade(grade)
    return gr.Dropdown(choices=subjects, value=subjects[0] if subjects else None)

# Build Gradio UI
with gr.Blocks(title="BOPPPS Lesson Plan Generator", theme=gr.themes.Soft()) as demo:
    
    gr.Markdown("""
    # BOPPPS Lesson Plan Generator
    ### Professional Instructional Design Tool
    Generate comprehensive lesson plans aligned with Pakistan National Curriculum 2006
    """)
    
    gr.Markdown("---")
    
    # Counter display
    with gr.Row():
        counter_display = gr.Number(
            label="Total Lesson Plans Generated",
            value=0,
            interactive=False,
            scale=1
        )
    
    gr.Markdown("---")
    
    with gr.Group():
        gr.Markdown("### School Information")
        school_name = gr.Textbox(
            label="School Name",
            placeholder="Enter the complete school name",
            info="This will appear as a header on the lesson plan"
        )
    
    with gr.Group():
        gr.Markdown("### Lesson Details")
        
        with gr.Row():
            grade = gr.Dropdown(
                choices=[f"Grade {i}" for i in range(1, 13)],
                label="Grade Level",
                value="Grade 9",
                info="Select the grade level for this lesson"
            )
            
            subject = gr.Dropdown(
                choices=get_subjects_for_grade("Grade 9"),
                label="Subject",
                value="English",
                info="Subject will update based on selected grade"
            )
        
        topic = gr.Textbox(
            label="Topic",
            placeholder="Enter the lesson topic",
            info="Be specific about the topic to be covered"
        )
        
        with gr.Row():
            students = gr.Number(
                label="Number of Students",
                value=30,
                minimum=1,
                maximum=200,
                step=1,
                info="Enter the exact number of students"
            )
            
            duration = gr.Dropdown(
                choices=["35", "40", "45", "50", "60", "70", "80"],
                label="Duration (minutes)",
                value="40",
                info="Select the class duration"
            )
        
        with gr.Row():
            date = gr.Textbox(
                label="Date",
                value=datetime.date.today().strftime("%Y-%m-%d"),
                placeholder="YYYY-MM-DD",
                info="Date of the lesson"
            )
            
            teacher = gr.Textbox(
                label="Teacher Name",
                placeholder="Enter teacher's full name",
                info="Name of the instructor"
            )
    
    gr.Markdown("---")
    
    with gr.Row():
        generate_button = gr.Button("Generate Lesson Plan", variant="primary", size="lg")
    
    output_file = gr.File(label="Download Lesson Plan Document")
    
    # Update subjects when grade changes
    grade.change(
        fn=update_subjects,
        inputs=[grade],
        outputs=[subject]
    )
    
    # Generate lesson plan
    generate_button.click(
        fn=generate_lesson_plan,
        inputs=[school_name, grade, subject, topic, students, duration, date, teacher],
        outputs=[output_file, counter_display]
    )
    
    gr.Markdown("---")
    gr.Markdown("""
    <div style='text-align: center; color: #666; font-size: 0.9em;'>
    <p>Developed by Najaf Ali Sharqi</p>
    <p>Based on Pakistan National Curriculum 2006</p>
    </div>
    """)

demo.launch()