app.py

import os
import streamlit as st
import requests
import json
import re
from docx import Document
from io import BytesIO
import plotly.express as px
import plotly.graph_objects as go
import pandas as pd
import networkx as nx
from collections import Counter

# --- Set Streamlit environment variables for HuggingFace compatibility ---
# These environment variables are an attempt to prevent permission errors
# by telling Streamlit to use a writable directory for its internal files.
# For full compatibility, you may also need a .streamlit/config.toml file
# in your repository with the content:
# [global]
# dataSavePath = "/tmp"
os.environ["STREAMLIT_BROWSER_GATHER_USAGE_STATS"] = "false"
os.environ["STREAMLIT_SERVER_ENABLE_ARROW_IPC"] = "false"
os.environ["STREAMLIT_SERVER_FOLDER"] = "/tmp"

# --- Page Configuration ---
st.set_page_config(
    page_title="Music Lesson Planner",
    page_icon="🎶",
    layout="wide",
    initial_sidebar_state="expanded"
)

# --- Constants and API Setup ---
# IMPORTANT: Set your Google API Key as an environment variable named GOOGLE_API_KEY
# You can get one from Google AI Studio: https://aistudio.google.com/app/apikey
GEMINI_API_KEY = os.getenv('GOOGLE_API_KEY')

# Base URL for Gemini API
GEMINI_API_BASE_URL = "https://generativelanguage.googleapis.com/v1beta/models/"

# Available Gemini Models for comparison
GEMINI_MODELS = {
    "Gemini 2.5 Flash": "gemini-2.5-flash",
    "Gemini 2.5 Pro": "gemini-2.5-pro",
    "Gemini 2.5 Flash Lite": "gemini-2.5-flash-lite",
}


# --- Helper Function for LLM API Call ---
def call_gemini_api(model_name, prompt_text, response_schema=None):
    """
    Calls the Gemini API with the given model and prompt.
    Handles JSON parsing and error reporting.
    """
    if not GEMINI_API_KEY:
        st.error(
            "Gemini API Key is not set. Please set the GOOGLE_API_KEY environment variable or replace `GEMINI_API_KEY = os.getenv('GOOGLE_API_KEY')` with your actual API key.")
        return None

    model_id = GEMINI_MODELS.get(model_name)
    if not model_id:
        st.error(f"Unknown model: {model_name}")
        return None

    url = f"{GEMINI_API_BASE_URL}{model_id}:generateContent?key={GEMINI_API_KEY}"

    headers = {
        "Content-Type": "application/json",
    }

    payload = {
        "contents": [
            {
                "role": "user",
                "parts": [{"text": prompt_text}]
            }
        ],
        "generationConfig": {}  # Initialize generationConfig
    }

    # If a response_schema is provided, configure for structured output
    if response_schema:
        payload["generationConfig"]["responseMimeType"] = "application/json"
        payload["generationConfig"]["responseSchema"] = response_schema

    try:
        response = requests.post(url, headers=headers, json=payload)
        response.raise_for_status()  # Raise an exception for HTTP errors (4xx or 5xx)
        response_data = response.json()

        if response_data and response_data.get("candidates"):
            # Access the text part of the response
            if response_schema:
                # For structured responses, the content is directly the JSON string
                raw_json_text = response_data["candidates"][0]["content"]["parts"][0]["text"]
                # Use regex to robustly extract the JSON object or array, ignoring any
                # surrounding text or malformed characters.
                json_match = re.search(r'\[.*\]|\{.*\}', raw_json_text, re.DOTALL)
                if json_match:
                    json_string = json_match.group(0)
                    try:
                        parsed_json = json.loads(json_string)
                        return parsed_json
                    except json.JSONDecodeError as e:
                        st.error(f"Failed to parse JSON for {model_name}. Error: {e}")
                        st.text_area(f"Raw output from {model_name}:", raw_json_text, height=200)
                        return None
                else:
                    st.error(f"Could not find a valid JSON object or array in the response from {model_name}.")
                    st.text_area(f"Raw output from {model_name}:", raw_json_text, height=200)
                    return None
            else:
                # For unstructured responses, return the text directly
                return response_data["candidates"][0]["content"]["parts"][0]["text"]
        else:
            st.error(f"No valid response candidates found for {model_name}.")
            st.json(response_data)  # Display the full response for debugging
            return None

    except requests.exceptions.HTTPError as http_err:
        st.error(f"HTTP error occurred for {model_name}: {http_err}")
        st.error(f"Response content: {response.text}")
        return None
    except requests.exceptions.ConnectionError as conn_err:
        st.error(f"Connection error occurred for {model_name}: {conn_err}")
        return None
    except requests.exceptions.Timeout as timeout_err:
        st.error(f"Timeout error occurred for {model_name}: {timeout_err}")
        return None
    except requests.exceptions.RequestException as req_err:
        st.error(f"An unexpected error occurred for {model_name}: {req_err}")
        return None
    except Exception as e:
        st.error(f"An unexpected error occurred during API call for {model_name}: {e}")
        return None


# --- Custom Markdown Formatting for Outline ---
def format_outline_for_display(outline_data, lesson_number):
    """
    Formats the outline JSON data for a single lesson into a human-readable markdown string.
    """
    if not outline_data:
        return "No outline data available."

    markdown_string = ""
    # Add Intro
    markdown_string += f"**Intro:**\n{outline_data.get('intro', 'N/A')}\n\n"

    # Add Key Teaching Points & Exercises
    markdown_string += "**Key Teaching Points & Exercise Suggestions:**\n"
    for point in outline_data.get('keyTeachingPoints', []):
        markdown_string += f"- **{point.get('point', 'N/A')}**\n"
        for exercise in point.get('exercises', []):
            markdown_string += f"  - {exercise}\n"

    # Add Outro
    markdown_string += f"\n**Outro:**\n{outline_data.get('outro', 'N/A')}"

    return markdown_string


# --- Visualization Functions ---
def analyze_lesson_complexity(lessons_data):
    """
    Analyzes lesson complexity based on key teaching points count and content length.
    Returns complexity scores for the complexity timeline.
    """
    complexity_scores = []
    for i, lesson in enumerate(lessons_data):
        # Base complexity on number of teaching points and content depth
        num_points = len(lesson.get('keyTeachingPoints', []))
        total_exercises = sum(len(point.get('exercises', [])) for point in lesson.get('keyTeachingPoints', []))

        # Simple scoring: more teaching points + more exercises = higher complexity
        complexity_score = num_points * 2 + total_exercises
        complexity_scores.append(complexity_score)

    return complexity_scores

def create_complexity_timeline(lessons_data):
    """Creates a complexity timeline visualization."""
    complexity_scores = analyze_lesson_complexity(lessons_data)

    # Create the timeline chart
    fig = go.Figure()

    # Add the complexity line
    fig.add_trace(go.Scatter(
        x=list(range(1, 6)),
        y=complexity_scores,
        mode='lines+markers',
        line=dict(color='#0277bd', width=3),
        marker=dict(size=12, color='#0277bd'),
        name='Complexity Score'
    ))

    # Add annotations for lesson types
    annotations = []
    for i, score in enumerate(complexity_scores):
        lesson_type = "Building" if i < 3 else "Reinforcing"
        annotations.append(dict(
            x=i+1, y=score,
            text=f"L{i+1}<br>{lesson_type}",
            showarrow=True,
            arrowhead=2,
            arrowsize=1,
            arrowwidth=2,
            arrowcolor='#666',
            bgcolor='white',
            bordercolor='#666',
            borderwidth=1
        ))

    fig.update_layout(
        title="Lesson Complexity Timeline",
        xaxis_title="Lesson Number",
        yaxis_title="Complexity Score",
        xaxis=dict(tickmode='linear', tick0=1, dtick=1),
        height=400,
        annotations=annotations
    )

    return fig

def extract_skills_from_lesson(lesson):
    """Extract key skills from a lesson's teaching points."""
    skills = []
    for point in lesson.get('keyTeachingPoints', []):
        skill = point.get('point', '').strip()
        if skill:
            skills.append(skill)
    return skills

def create_skill_flow_diagram(lessons_data):
    """Creates a skill building flow diagram using networkx and plotly."""
    # Create a directed graph
    G = nx.DiGraph()

    # Add nodes for each lesson and extract skills
    lesson_skills = {}
    all_skills = []

    for i, lesson in enumerate(lessons_data):
        lesson_name = f"Lesson {i+1}"
        skills = extract_skills_from_lesson(lesson)
        lesson_skills[lesson_name] = skills
        all_skills.extend(skills)

    # Add lesson nodes
    for i in range(5):
        lesson_name = f"Lesson {i+1}"
        G.add_node(lesson_name, node_type='lesson', lesson_num=i+1)

    # Add skill dependencies (lessons flow into next lesson)
    for i in range(4):
        G.add_edge(f"Lesson {i+1}", f"Lesson {i+2}")

    # Create layout
    pos = nx.spring_layout(G, k=3, iterations=50)

    # Extract node and edge information for plotly
    edge_x, edge_y = [], []
    for edge in G.edges():
        x0, y0 = pos[edge[0]]
        x1, y1 = pos[edge[1]]
        edge_x.extend([x0, x1, None])
        edge_y.extend([y0, y1, None])

    node_x = [pos[node][0] for node in G.nodes()]
    node_y = [pos[node][1] for node in G.nodes()]
    node_text = [f"{node}<br>Skills: {len(lesson_skills[node])}" for node in G.nodes()]

    # Create the figure
    fig = go.Figure()

    # Add edges
    fig.add_trace(go.Scatter(
        x=edge_x, y=edge_y,
        line=dict(width=2, color='#666'),
        hoverinfo='none',
        mode='lines',
        showlegend=False
    ))

    # Add nodes
    fig.add_trace(go.Scatter(
        x=node_x, y=node_y,
        mode='markers+text',
        marker=dict(size=[20 + len(lesson_skills[node])*3 for node in G.nodes()],
                   color=['#0277bd', '#0288d1', '#039be5', '#03a9f4', '#29b6f6'],
                   line=dict(width=2, color='white')),
        text=[node.replace(' ', '<br>') for node in G.nodes()],
        textposition='middle center',
        textfont=dict(size=10, color='white'),
        hovertext=node_text,
        hoverinfo='text',
        showlegend=False
    ))

    fig.update_layout(
        title="Skill Building Flow",
        showlegend=False,
        xaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
        yaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
        height=400,
        margin=dict(l=20, r=20, t=40, b=20)
    )

    return fig

def create_topic_coverage_heatmap(lessons_data):
    """Creates a heatmap showing topic coverage across lessons."""
    # Extract all unique topics/concepts from teaching points
    all_topics = []
    lesson_topics = {}

    # Common music theory categories to look for
    topic_categories = {
        'scales': ['scale', 'major', 'minor', 'chromatic'],
        'rhythm': ['beat', 'tempo', 'rhythm', 'timing', 'bpm'],
        'technique': ['finger', 'hand', 'practice', 'exercise'],
        'theory': ['chord', 'interval', 'sharp', 'flat', 'key'],
        'performance': ['play', 'sound', 'listen', 'hear']
    }

    # Analyze each lesson
    coverage_matrix = []
    for i, lesson in enumerate(lessons_data):
        lesson_coverage = {}
        lesson_text = json.dumps(lesson).lower()

        for category, keywords in topic_categories.items():
            coverage_score = sum(lesson_text.count(keyword) for keyword in keywords)
            lesson_coverage[category] = coverage_score

        coverage_matrix.append(lesson_coverage)

    # Create DataFrame
    df = pd.DataFrame(coverage_matrix, index=[f'Lesson {i+1}' for i in range(5)])

    # Create heatmap
    fig = px.imshow(
        df.T,
        labels=dict(x="Lesson", y="Topic Category", color="Coverage Intensity"),
        title="Topic Coverage Heatmap",
        color_continuous_scale="Blues",
        aspect="auto"
    )

    fig.update_layout(height=400)

    return fig


# --- DOCX Conversion Function ---
def strip_html_tags(text):
    """
    Remove HTML tags from text for DOCX export.
    """
    if not text:
        return text
    # Remove span tags with class attributes
    text = re.sub(r'<span class=["\'][^"\']*["\']>', '', text)
    text = re.sub(r'</span>', '', text)
    return text

def sanitize_filename(text, max_length=100):
    """
    Sanitizes text to create a safe filename for both filesystems and HTTP headers.

    Args:
        text: The text to sanitize
        max_length: Maximum length for the sanitized text (default 100)

    Returns:
        A sanitized string safe for use in filenames and HTTP Content-Disposition headers
    """
    if not text:
        return "lesson"

    # Replace newlines, tabs, and other whitespace with single space
    text = re.sub(r'[\n\r\t\v\f]+', ' ', text)

    # Replace multiple spaces with single space
    text = re.sub(r'\s+', ' ', text)

    # Remove or replace unsafe characters for filenames and HTTP headers
    # Keep only alphanumeric, spaces, hyphens, and underscores
    text = re.sub(r'[^\w\s\-]', '', text)

    # Replace spaces with underscores
    text = text.replace(' ', '_')

    # Remove leading/trailing underscores
    text = text.strip('_')

    # Truncate to max_length while avoiding cutting mid-word
    if len(text) > max_length:
        text = text[:max_length].rsplit('_', 1)[0]

    # Ensure we have at least some text
    if not text:
        return "lesson"

    return text

def create_docx_file(lessons_data, lesson_topic, lesson_length, model_name):
    """
    Creates a DOCX file from a sequence of lessons.
    """
    document = Document()

    document.add_heading(f"Lesson Plan Sequence: {lesson_topic}", level=1)
    document.add_paragraph(f"Length per lesson: {lesson_length}")
    document.add_paragraph(f"Generated by: {model_name}")
    document.add_paragraph("\n")

    for i in range(5):
        outline_data = lessons_data['outlines'][i]
        draft_text = lessons_data['drafts'][i]

        document.add_heading(f"Lesson {i + 1}", level=2)
        document.add_paragraph("\n")

        # Add Outline Sections
        document.add_heading("Outline", level=3)
        document.add_paragraph(f"**Intro:**\n{strip_html_tags(outline_data.get('intro', 'N/A'))}")

        document.add_heading("Key Teaching Points & Exercise Suggestions", level=4)
        for point in outline_data.get('keyTeachingPoints', []):
            document.add_paragraph(f"- {strip_html_tags(point.get('point', 'N/A'))}", style='List Bullet')
            for exercise in point.get('exercises', []):
                document.add_paragraph(strip_html_tags(exercise), style='List Bullet')

        document.add_heading("Outro", level=4)
        document.add_paragraph(strip_html_tags(outline_data.get('outro', 'N/A')))

        # Add Full Draft Content (strip HTML tags)
        document.add_heading("Full Lesson Draft", level=2)
        document.add_paragraph(strip_html_tags(draft_text))
        document.add_paragraph("\n")

    byte_io = BytesIO()
    document.save(byte_io)
    byte_io.seek(0)
    return byte_io.getvalue()


# --- Password Protection ---
def authenticate_user():
    st.markdown("## 🔐 Secure Login")
    password = st.text_input("Password", type="password", key="password_input")
    submit = st.button("Login", key="login_button")

    if submit:
        correct_password = os.getenv("APP_PASSWORD")

        if password == correct_password:
            st.session_state["authenticated"] = True
            st.rerun()  # Rerun to clear password input and show app content
        else:
            st.error("Invalid password")


# Check authentication status
if "authenticated" not in st.session_state:
    st.session_state["authenticated"] = False

if not st.session_state["authenticated"]:
    authenticate_user()
    st.stop()  # Stop execution if not authenticated
else:
    # --- Main Application Logic (Protected by Authentication) ---
    st.title("🎶 Music Lesson Planner")

    st.markdown("""
    This app helps you draft outlines and detailed lesson plans for online music lessons using different Gemini models.
    Compare outputs to find the best fit for your pedagogical needs!
    """)

    # Initialize session state for outlines and drafts if not already present
    if 'lessons_data' not in st.session_state:
        st.session_state.lessons_data = {}

    # Input fields for lesson details
    with st.sidebar:
        st.header("Lesson Details")
        lesson_topic = st.text_area("Lesson Topic", "Introduction to Solfege", height=100)
        lesson_length = st.selectbox("Lesson Length", ["5-minute", "7-minute", "10-minute"], index=0)

        st.header("Model Selection")
        selected_models = st.multiselect(
            "Select Gemini Models",
            list(GEMINI_MODELS.keys()),
            default=["Gemini 2.5 Pro", "Gemini 2.5 Flash"]
        )

        st.header("Prompt Customization")
        default_outline_system_prompt = (
            "You are an AI assistant specialized in creating concise and structured outlines for online music lessons. "
            "Your goal is to provide a clear, pedagogical framework that music educators can easily follow. "
            "Focus on three main sections: an introduction, a list of key teaching points with suggested exercises, and a conclusion (outro)."
            "The lessons are online and asynchronous, so ensure the content is suitable for self-paced learning."
            "DO NOT include any quizzes, assessments, images, or audio/video."
        )
        outline_system_prompt = st.text_area("Outline System Prompt", default_outline_system_prompt, height=150)

        # Static user prompt template for the outline
        outline_user_prompt_template = (
            "Create a sequence of five online music lessons on the topic of '{lesson_topic}'. "
            "Each lesson in the sequence should be a {lesson_length} lesson. "
            "The sequence should build in complexity from lessons 1 through 3, with lessons 4 and 5 focusing on reinforcement and review. "
            "The entire sequence is a skill pack: a structured sequence of 5 short videos, designed to teach a specific skill through guided, real-time repetition."
            "Ideally, a student will view one video from the skill pack per day. This gives them some time in between videos, rather than bingeing all 5 at once."
            "This is a play-to-learn format, with the student playing along with the instructor for the entire session. This requires the instructor to coach the viewer along while simultaneously playing through the examples."
            "All practice is done in real time, ensuring learners can build muscle memory while they play—no extra practice required. As a result, the increase in difficulty from video to video should be very minimal."

            "The entire response must be a JSON array containing five lesson objects. "
            "Each lesson object must contain the following three keys: 'intro' (a single paragraph string), 'keyTeachingPoints' (a list of objects, each with 'point' and 'exercises'), and 'outro' (a single paragraph string)."
            "Please provide the response strictly in JSON format according to the schema provided, with no additional text or markdown outside the JSON array."
        )

        # Sample script to provide tone reference for draft generation (not displayed to the user)
        sample_script_text = """Hey everyone, [instructor name] here, community manager for Piano and welcome to Unit 2 of the Piano Curriculum. Hope you've had plenty of time to practice and integrate the concepts and exercises and scales that we learned in that first unit. Got those chords and different majors, scale sounds under your hands. First unit, we're going to dive a little further into learning some new keys, some new ways to shape chords. And to help me do that, I have my friend [instructor name] here. Hello, how are you doing? I'm doing well. Great. It's good to be here. Tell us a little bit about yourself, how long you've been playing piano, how long you've been teaching for. I've been playing piano since I was about nine. I begged and begged for piano lessons and in retrospect I might have gotten myself into a bit of a disaster. But I'm very thankful that my parents have helped me to stick to it. So I've been playing from about nine and I've been teaching for about ten years. Nice. Yeah. So a lot of experience as a teacher. Lisa is going to teach us a little bit about the G major scale today, which is a whole new key. Lisa, tell us a little bit about G major. Okay. G major is, it's a really cool scale because it introduces us to our first sharp. So you're probably wondering what is a sharp. You've seen them before. They look like a hashtag. They are the original hashtag. And what a sharp does is it raises the tone by a half step. So what I mean by that is if we were to start on C, a sharp would mean that we're going to play the C sharp here. So it simply raises a note by a half step or a half tone. In our scale today, we're going to need to play an F sharp. So I'm going to show you where that is. So C, D, E, F. To play an F sharp, we go here. Let's F sharp. So the reason why it's important to know what a sharp is, is it because it allows us to follow the major scale tone pattern that we need in order to make the major scale sound like a major scale. So if we were to play a G scale starting on G and just play all the white keys, it would sound like this. Which doesn't, it doesn't sound quite right. So let's take a look at that step by step and we'll see how following that formula for a major scale introduces the sharp. So if we go from G to A, that's the whole step. From A to B, whole step. B to C, there's our half step. C to D, whole step. D to E, whole step. Now we need another whole step here. So in order to make that happen, we have to move to our F sharp. And then the scale ends nicely on the half step. So that's your G scale. I'm going to play it once more for you so you can hear it all together. And so we go back down the same way we came. And the coordination between our fourth finger and our third finger can be a little bit tricky at first. So you have to be careful that your four lands on the F sharp and your three is going to hit the E and we're going to come all the way down like this. And there's your G scale. Awesome. So in order to play anything in the key of G, we have to have that F sharp. Yes. So if you're playing in the key of G and an F comes around, it's going to be a sharp unless otherwise stated. Right. And how does that scale look in the left hand? Looks just like this. We're going to start with our five finger on G and we're going to go five, four, three, two, one. Fly over with our three. Remember our F sharp and then we've got our G. And we go down the same way we went up. Awesome. So it has the same fingering pattern that you guys learned playing the C major scale. Yes. Keep an eye out for that F sharp, the seventh note in the scale. So how do you recommend practicing the G major scale? So I recommend practicing the G major scale first with your right hand. It's up to you. You can start right hand or left hand. I'm definitely more dominant in my right. So it's easier. You want to make it easier on yourself when you start. So I'd suggest that you play the scale a bunch of times with your right hand, maybe five, bunch of times with your left hand and then we're going to try hands together. So hands together can be tricky at first. Don't get frustrated. It takes practice. You will get it. Go very slowly. And I always suggest letting yourself put weight in the keys. So if you're really struggling, slow down and be intentional. Pretend like your hands are heavy. And that's going to help your mind connect to your hand muscles so that you can develop the muscle memory you need for the scale. And that's how it sounds hands together. Awesome. So the way that I recommend practicing the G major scale, play it at a tempo, at a nice slow tempo that's comfortable for you. Try about 60 BPM. It'll sound really slow, but it's really good for building that muscle memory and just the basic theory knowledge. Go up and down the scale one octave, five times. If you can do that without making any mistakes, if you can keep yourself a solid sense of rhythm as you're playing, you'll be ready to move on to the next video. Lisa, can you demonstrate that, this scale, with a quick? Yes, absolutely. Okay. Perfect. So practice that G major scale up and down five times at a slow tempo. When you can do that without any mistakes and keep a stable rhythm, you'll be ready to move on to the next video. If you have any questions or need any clarification, let me know with an email to Jordan at Piano.com. I'll be happy to help you out and we'll see you at the next video."""

        draft_system_prompt_base = (
            "You are an AI assistant specialized in expanding structured lesson outlines into detailed, engaging rough drafts for online music lessons. "
            "Your goal is to provide specific examples and pedagogical details for each teaching point and exercise. "
            "The language should be engaging and professional, tailored for music educators."
            "The lessons are online and asynchronous, so ensure the draft is suitable for self-paced learning."
            "DO NOT include any quizzes, assessments, images, or audio/video."
            "\n\n"
            "IMPORTANT: Apply automatic color coding by wrapping relevant terms in HTML spans with specific CSS classes:\n"
            "- Music theory terms (scales, intervals, chords, keys, notes, etc.): <span class='music-theory'>term</span>\n"
            "- Tempo & timing terms (BPM, rhythm, beat, tempo, etc.): <span class='tempo-timing'>term</span>\n"
            "- Practice instructions (practice, repeat, try, work on, etc.): <span class='practice-instruction'>term</span>\n"
            "- Reminders and callbacks (last time we did, in the previous lesson, last time, earlier we learned, etc.): <span class='student-engagement'>phrase</span>\n"
            "Apply this formatting naturally throughout your response without changing the content or flow."
        )
        draft_system_prompt = st.text_area("Draft System Prompt", draft_system_prompt_base, height=200)

        generate_button = st.button("Generate Lesson Plans")

        # Add a logout button to the sidebar
        if st.session_state["authenticated"]:
            if st.button("Logout", key="logout_button_sidebar"):
                st.session_state["authenticated"] = False
                st.rerun()

    # --- Define Outline Schema (Revised) ---
    outline_response_schema = {
        "type": "ARRAY",
        "items": {
            "type": "OBJECT",
            "properties": {
                "intro": {"type": "STRING"},
                "keyTeachingPoints": {
                    "type": "ARRAY",
                    "items": {
                        "type": "OBJECT",
                        "properties": {
                            "point": {"type": "STRING"},
                            "exercises": {"type": "ARRAY", "items": {"type": "STRING"}}
                        },
                        "required": ["point", "exercises"]
                    }
                },
                "outro": {"type": "STRING"}
            },
            "required": ["intro", "keyTeachingPoints", "outro"]
        }
    }

    # --- Lesson Generation Logic (Triggered by button) ---
    if generate_button:
        # Clear previous results when new generation is triggered
        st.session_state.lessons_data = {}
        st.session_state.lesson_topic = lesson_topic  # Store for download
        st.session_state.lesson_length = lesson_length  # Store for download
        st.session_state.selected_models = selected_models  # Store for download

        # Generate outlines
        for model_name in selected_models:
            current_outline_user_prompt = outline_user_prompt_template.format(
                lesson_length=lesson_length,
                lesson_topic=lesson_topic
            )
            full_outline_prompt = f"{outline_system_prompt}\n{current_outline_user_prompt}"
            all_outlines = call_gemini_api(model_name, full_outline_prompt, outline_response_schema)
            if all_outlines and len(all_outlines) == 5:
                st.session_state.lessons_data[model_name] = {'outlines': all_outlines, 'drafts': []}
                # Generate drafts
                for i, outline_data in enumerate(all_outlines):
                    outline_for_draft = json.dumps(outline_data, indent=2)

                    # Combine the editable prompt with the hidden sample script
                    full_draft_system_prompt = (
                        f"{draft_system_prompt}\n\n"
                        "Use the following example script as a reference for tone and style: \n\n"
                        f"--- START OF SAMPLE SCRIPT ---\n{sample_script_text}\n--- END OF SAMPLE SCRIPT ---"
                    )

                    draft_prompt = (
                        f"{full_draft_system_prompt}\n\n"
                        f"Expand the following outline for Lesson {i + 1} into a detailed rough draft for a {lesson_length} lesson. "
                        "Ensure the language is engaging for music educators.\n\n"
                        f"Outline:\n```json\n{outline_for_draft}\n```"
                    )

                    raw_draft_text = call_gemini_api(model_name, draft_prompt)
                    if raw_draft_text:
                        st.session_state.lessons_data[model_name]['drafts'].append(raw_draft_text)
                    else:
                        st.session_state.lessons_data[model_name]['drafts'].append(None)
            else:
                st.error(f"Failed to generate a complete 5-lesson sequence for {model_name}. Please try again.")

    # --- Display Generated Content and Download Buttons (Always displayed if in session_state) ---
    if st.session_state.get('lessons_data'):
        for model_name, lessons in st.session_state.lessons_data.items():
            st.subheader(f"Lesson Plans from {model_name}")

            for i in range(5):
                outline_data = lessons['outlines'][i] if len(lessons['outlines']) > i else None
                draft_text = lessons['drafts'][i] if len(lessons['drafts']) > i else None

                expander_title = f"Lesson {i + 1}: "
                if outline_data and 'intro' in outline_data:
                    expander_title += outline_data['intro'][:50] + "..."
                else:
                    expander_title += "Outline could not be generated."

                with st.expander(expander_title):
                    outline_col, draft_col = st.columns(2)
                    with outline_col:
                        st.markdown(f"**Outline for Lesson {i + 1}**")
                        if outline_data:
                            human_readable_outline = format_outline_for_display(outline_data, i + 1)
                            st.markdown(f'<div class="result-box">{human_readable_outline}</div>',
                                        unsafe_allow_html=True)
                        else:
                            st.markdown(
                                f'<div class="result-box text-gray-500">Could not generate outline for this lesson.</div>',
                                unsafe_allow_html=True)

                    with draft_col:
                        st.markdown(f"**Rough Draft for Lesson {i + 1}**")
                        if draft_text:
                            st.markdown(f'<div class="result-box">{draft_text}</div>', unsafe_allow_html=True)
                        else:
                            st.markdown(
                                f'<div class="result-box text-gray-500">Could not generate draft for this lesson.</div>',
                                unsafe_allow_html=True)

        # --- Visualizations Section ---
        st.subheader("📊 Lesson Analysis & Visualizations")

        # Show visualizations for each model
        for model_name, lessons in st.session_state.lessons_data.items():
            if len(lessons['outlines']) == 5:  # Only show if we have complete data
                st.write(f"**Analysis for {model_name}:**")

                try:
                    heatmap_fig = create_topic_coverage_heatmap(lessons['outlines'])
                    st.plotly_chart(heatmap_fig, use_container_width=True)
                    st.write("This heatmap shows which music theory topics are emphasized in each lesson.")
                except Exception as e:
                    st.error(f"Error generating topic coverage heatmap: {e}")

                st.divider()  # Add visual separation between models

        st.subheader("Download All Lessons")
        download_cols = st.columns(len(st.session_state.get('selected_models', [])))
        for i, model_name in enumerate(st.session_state.get('selected_models', [])):
            with download_cols[i]:
                lessons_data = st.session_state.lessons_data.get(model_name)
                if lessons_data:
                    num_outlines = len(lessons_data.get('outlines', []))
                    num_drafts = len(lessons_data.get('drafts', []))

                    if num_outlines == 5 and num_drafts == 5:
                        docx_file = create_docx_file(
                            lessons_data,
                            st.session_state.lesson_topic,
                            st.session_state.lesson_length,
                            model_name
                        )
                        # Sanitize filename components
                        safe_model_name = sanitize_filename(model_name, max_length=50)
                        safe_topic = sanitize_filename(st.session_state.lesson_topic, max_length=80)

                        st.download_button(
                            label=f"Download {model_name} (DOCX)",
                            data=docx_file,
                            file_name=f"{safe_model_name}_lesson_sequence_{safe_topic}.docx",
                            mime="application/vnd.openxmlformats-officedocument.wordprocessingml.document",
                            key=f"download_docx_{model_name}"
                        )
                    else:
                        st.markdown(
                            f'<div class="text-gray-500">Cannot download DOCX for {model_name}<br>Outlines: {num_outlines}/5, Drafts: {num_drafts}/5</div>',
                            unsafe_allow_html=True)
                else:
                    st.markdown(
                        f'<div class="text-gray-500">No data for {model_name}</div>',
                        unsafe_allow_html=True)
    elif generate_button:
        st.info("No content generated. Please check for API errors or adjust prompts.")

    st.markdown("""
    <style>
    .result-box {
        border: 1px solid #ddd;
        border-radius: 8px;
        padding: 15px;
        margin-bottom: 10px;
        background-color: #f9f9f9;
        white-space: pre-wrap;
        word-wrap: break-word;
    }
    .text-gray-500 {
        color: #6b7280;
    }

    /* Color coding for music lesson terms */
    .music-theory {
        background-color: #e0f2fe;
        color: #0277bd;
        padding: 1px 3px;
        border-radius: 3px;
        font-weight: 500;
    }
    .tempo-timing {
        background-color: #f3e5f5;
        color: #7b1fa2;
        padding: 1px 3px;
        border-radius: 3px;
        font-weight: 500;
    }
    .practice-instruction {
        background-color: #fff3e0;
        color: #f57c00;
        padding: 1px 3px;
        border-radius: 3px;
        font-weight: 500;
    }
    .student-engagement {
        background-color: #e8f5e8;
        color: #2e7d32;
        padding: 1px 3px;
        border-radius: 3px;
        font-weight: 500;
    }
    </style>
    """, unsafe_allow_html=True)