import streamlit as st
from gpt_method import (
    generate_lesson_plan,
    generate_objectives,
    generate_materials,
    generate_steps,
    generate_assessment
)

def edit_lesson_plan():
    st.title("📝 Edit Lesson Plan")
    
    # Load existing lesson plan data from session state
    if 'lesson_plan' not in st.session_state:
        st.error("No lesson plan found. Please create a lesson plan first.")
        return
    
    # Create tabs for different sections
    tab1, tab2, tab3, tab4, tab5 = st.tabs([
        "Basic Info", "Objectives", "Materials", "Steps", "Assessment"
    ])
    
    with tab1:
        st.subheader("Basic Information")
        lesson_name = st.text_input(
            "Course Name",
            value=st.session_state.lesson_plan.get('lesson_name', '')
        )
        lesson_theme = st.text_input(
            "Course Theme",
            value=st.session_state.lesson_plan.get('lesson_theme', '')
        )
        grade_level = st.selectbox(
            "Student Level",
            ["Elementary School", "Middle School", "High School", "College", "Graduate"],
            index=["Elementary School", "Middle School", "High School", "College", "Graduate"].index(
                st.session_state.lesson_plan.get('grade_level', 'College')
            )
        )
        duration = st.number_input(
            "Duration (minutes)",
            min_value=15,
            step=15,
            value=st.session_state.lesson_plan.get('duration', 60)
        )
    
    with tab2:
        st.subheader("Learning Objectives")
        objectives = st.text_area(
            "Course Objectives",
            value=st.session_state.lesson_plan.get('objectives', ''),
            height=200
        )
        if st.button("Regenerate Objectives"):
            with st.spinner("Generating new objectives..."):
                objectives = generate_objectives(
                    lesson_theme,
                    lesson_name=lesson_name,
                    lesson_type=grade_level,
                    existing_goals=objectives
                )
                st.session_state.lesson_plan['objectives'] = objectives
                st.experimental_rerun()
    
    with tab3:
        st.subheader("Materials")
        materials = st.text_area(
            "Required Materials",
            value=st.session_state.lesson_plan.get('materials', ''),
            height=200
        )
        if st.button("Regenerate Materials"):
            with st.spinner("Generating materials list..."):
                materials = generate_materials(lesson_theme, grade_level)
                st.session_state.lesson_plan['materials'] = materials
                st.experimental_rerun()
    
    with tab4:
        st.subheader("Lesson Steps")
        steps = st.text_area(
            "Lesson Steps",
            value=st.session_state.lesson_plan.get('steps', ''),
            height=300
        )
        if st.button("Regenerate Steps"):
            with st.spinner("Generating lesson steps..."):
                steps = generate_steps(lesson_theme, grade_level, duration)
                st.session_state.lesson_plan['steps'] = steps
                st.experimental_rerun()
    
    with tab5:
        st.subheader("Assessment")
        assessment = st.text_area(
            "Assessment Methods",
            value=st.session_state.lesson_plan.get('assessment', ''),
            height=200
        )
        if st.button("Regenerate Assessment"):
            with st.spinner("Generating assessment methods..."):
                assessment = generate_assessment(lesson_theme, grade_level)
                st.session_state.lesson_plan['assessment'] = assessment
                st.experimental_rerun()
    
    # Save changes button
    if st.button("Save Changes"):
        # Update session state with new values
        st.session_state.lesson_plan.update({
            'lesson_name': lesson_name,
            'lesson_theme': lesson_theme,
            'grade_level': grade_level,
            'duration': duration,
            'objectives': objectives,
            'materials': materials,
            'steps': steps,
            'assessment': assessment
        })
        
        # Generate updated lesson plan
        updated_plan = generate_lesson_plan(
            lesson_name,
            lesson_theme,
            grade_level,
            objectives,
            materials,
            duration,
            steps,
            assessment
        )
        
        st.session_state.lesson_plan['full_plan'] = updated_plan
        st.success("Changes saved successfully!")
        
        # Show preview of updated plan
        st.subheader("Updated Lesson Plan Preview")
        st.markdown(updated_plan)
        
        # Add download button
        st.download_button(
            label="Download Updated Lesson Plan",
            data=updated_plan,
            file_name=f"lesson_plan_{lesson_name.lower().replace(' ', '_')}.txt",
            mime="text/plain"
        )