app.py

import streamlit as st
import google.generativeai as genai
import os
import json
import base64
from dotenv import load_dotenv
from streamlit_local_storage import LocalStorage
import re
import streamlit.components.v1 as components
import math # Needed for trigonometry in dynamic visuals

# --- PAGE CONFIGURATION ---
st.set_page_config(
    page_title="Math Jegna - Your AI Math Tutor",
    page_icon="🧠",
    layout="wide"
)

# Create an instance of the LocalStorage class
localS = LocalStorage()

# --- HELPER FUNCTIONS ---
def format_chat_for_download(chat_history):
    """Formats the chat history into a human-readable string for download."""
    formatted_text = f"# Math Mentor Chat\n\n"
    for message in chat_history:
        role = "You" if message["role"] == "user" else "Math Mentor"
        formatted_text += f"**{role}:**\n{message['content']}\n\n---\n\n"
    return formatted_text

def convert_role_for_gemini(role):
    """Convert Streamlit chat roles to Gemini API roles"""
    if role == "assistant":
        return "model"
    return role  # "user" stays the same

def should_generate_visual(user_prompt, ai_response):
    """Determine if a visual aid would be helpful based on the content"""
    # Expanded keywords to trigger new dynamic visuals
    k12_visual_keywords = [
        'add', 'subtract', 'multiply', 'times', 'divide', 'divided by', 'counting', 'numbers',
        'fraction', 'half', 'quarter', 'third', 'parts', 'whole',
        'shape', 'triangle', 'circle', 'square', 'rectangle',
        'money', 'coins', 'dollars', 'cents', 'change',
        'time', 'clock', 'hours', 'minutes', 'o\'clock',
        'measurement', 'length', 'height', 'weight',
        'place value', 'tens', 'ones', 'hundreds',
        'pattern', 'sequence', 'skip counting',
        'greater than', 'less than', 'equal', 'compare',
        'number line', 'array', 'grid', 'area model'
    ]
    
    combined_text = (user_prompt + " " + ai_response).lower()
    # Also check for symbols
    return any(keyword in combined_text for keyword in k12_visual_keywords) or any(op in user_prompt for op in ['*', '/'])


def create_visual_manipulative(user_prompt, ai_response):
    """-- SMART VISUAL ROUTER --
    Parses the user prompt and calls the appropriate dynamic visual function."""
    try:
        user_lower = user_prompt.lower().replace(' ', '') # Remove spaces for easier parsing
        
        # Priority 1: Division (e.g., "42 divided by 6", "21 / 3")
        div_match = re.search(r'(\d+)dividedby(\d+)', user_lower) or re.search(r'(\d+)/(\d+)', user_lower)
        if div_match and "fraction" not in user_lower: # Avoid confusion with fractions
            dividend, divisor = int(div_match.group(1)), int(div_match.group(2))
            if dividend <= 50 and divisor > 0: # Keep visuals manageable
                return create_division_groups_visual(dividend, divisor)

        # Priority 2: Multiplication (e.g., "3 times 5", "15 * 19")
        mult_match = re.search(r'(\d+)(?:x|times|\*)(\d+)', user_lower)
        if mult_match:
            num1, num2 = int(mult_match.group(1)), int(mult_match.group(2))
            # Use dot array for small numbers
            if num1 <= 10 and num2 <= 10:
                return create_multiplication_array(num1, num2)
            # Use area model for larger numbers (up to 99x99)
            elif 10 < num1 < 100 and 10 < num2 < 100:
                return create_multiplication_area_model(num1, num2)

        # Priority 3: Time / Clock (e.g., "7:30", "4 o'clock")
        time_match = re.search(r'(\d{1,2}):(\d{2})', user_lower) or re.search(r'(\d{1,2})o\'clock', user_lower)
        if time_match:
            groups = time_match.groups()
            hour = int(groups[0])
            minute = int(groups[1]) if len(groups) > 1 and groups[1] else 0
            if 1 <= hour <= 12 and 0 <= minute <= 59:
                return create_clock_visual(hour, minute)

        # Priority 4: Fractions (e.g., "2/5", "fraction 3/8")
        fraction_match = re.search(r'(\d+)/(\d+)', user_lower)
        if fraction_match:
            num, den = int(fraction_match.group(1)), int(fraction_match.group(2))
            if 0 < num <= den and den <= 16:
                 return create_dynamic_fraction_circle(num, den)

        # Priority 5: Addition/Subtraction Blocks
        if any(word in user_lower for word in ['add', 'plus', '+', 'subtract', 'minus', 'takeaway', '-']):
            numbers = re.findall(r'\d+', user_prompt)
            if len(numbers) >= 2:
                num1, num2 = int(numbers[0]), int(numbers[1])
                operation = 'add' if any(w in user_lower for w in ['add', 'plus', '+']) else 'subtract'
                if num1 <= 20 and num2 <= 20:
                    return create_counting_blocks(num1, num2, operation)

        # Priority 6: Number Lines
        if 'numberline' in user_lower:
            numbers = [int(n) for n in re.findall(r'\d+', user_prompt)]
            if numbers:
                start = min(numbers) - 2
                end = max(numbers) + 2
                return create_number_line(start, end, numbers, "Your Numbers on the Line")

        # Priority 7: Place Value
        if 'placevalue' in user_lower:
            numbers = re.findall(r'\d+', user_prompt)
            if numbers:
                num = int(numbers[0])
                if num <= 999:
                    return create_place_value_blocks(num)
        
        # Fallback to static, general visuals
        if any(word in user_lower for word in ['fraction', 'part']): return create_dynamic_fraction_circle(1, 2)
        if any(word in user_lower for word in ['shape']): return create_shape_explorer()
        if any(word in user_lower for word in ['money', 'coin']): return create_money_counter()
        if any(word in user_lower for word in ['time', 'clock']): return create_clock_visual(10, 10)

        return None # No relevant visual found
        
    except Exception as e:
        st.error(f"Could not create visual: {e}")
        return None

# --- VISUAL TOOLBOX FUNCTIONS ---

def create_multiplication_area_model(num1, num2):
    """(NEW & Dynamic) Creates an area model for 2-digit by 2-digit multiplication."""
    n1_tens, n1_ones = num1 // 10, num1 % 10
    n2_tens, n2_ones = num2 // 10, num2 % 10

    p1 = n1_tens * n2_tens * 100  # E.g., 10 * 10
    p2 = n1_tens * n2_ones * 10   # E.g., 10 * 9
    p3 = n1_ones * n2_tens * 10   # E.g., 5 * 10
    p4 = n1_ones * n2_ones       # E.g., 5 * 9
    
    total = p1 + p2 + p3 + p4

    html = f"""
    <div style="font-family: sans-serif; padding: 20px; background: linear-gradient(135deg, #f5f7fa 0%, #c3cfe2 100%); border-radius: 15px; margin: 10px 0;">
        <h3 style="text-align: center; color: #333;">Area Model for {num1} × {num2}</h3>
        <div style="display: grid; grid-template-columns: auto 1fr 1fr; grid-template-rows: auto 1fr 1fr; gap: 5px; max-width: 400px; margin: 20px auto; color: #333;">
            <!-- Corner --> <div/>
            <div style="text-align: center; font-weight: bold; padding: 5px;">{n2_tens*10}</div>
            <div style="text-align: center; font-weight: bold; padding: 5px;">{n2_ones}</div>
            
            <div style="text-align: center; font-weight: bold; padding: 5px;">{n1_tens*10}</div>
            <div style="background: #FFADAD; padding: 20px; text-align: center; border-radius: 8px;"><b>{n1_tens*10}×{n2_tens*10}</b><br/>{p1}</div>
            <div style="background: #FFD6A5; padding: 20px; text-align: center; border-radius: 8px;"><b>{n1_tens*10}×{n2_ones}</b><br/>{p2}</div>
            
            <div style="text-align: center; font-weight: bold; padding: 5px;">{n1_ones}</div>
            <div style="background: #FDFFB6; padding: 20px; text-align: center; border-radius: 8px;"><b>{n1_ones}×{n2_tens*10}</b><br/>{p3}</div>
            <div style="background: #CAFFBF; padding: 20px; text-align: center; border-radius: 8px;"><b>{n1_ones}×{n2_ones}</b><br/>{p4}</div>
        </div>
        <div style="text-align: center; margin-top: 15px; font-size: 1.1em;">
            <b>Add the parts together:</b> {p1} + {p2} + {p3} + {p4} = <b>{total}</b>
        </div>
    </div>
    """
    return html

def create_division_groups_visual(dividend, divisor):
    """(NEW & Dynamic) Creates a visual for division by grouping."""
    if divisor == 0: return ""
    quotient = dividend // divisor
    
    groups_html = ""
    dot_colors = ["#FF6B6B", "#4ECDC4", "#FFD93D", "#95E1D3", "#A0C4FF", "#FDBF6F"]
    
    for i in range(divisor):
        dots_in_group = "".join([f'<div style="width: 15px; height: 15px; background: {dot_colors[i % len(dot_colors)]}; border-radius: 50%;"></div>' for _ in range(quotient)])
        groups_html += f"""
            <div style="border: 2px dashed {dot_colors[i % len(dot_colors)]}; border-radius: 10px; padding: 10px; text-align: center;">
                <b style="color: #333;">Group {i+1}</b>
                <div style="display: flex; flex-wrap: wrap; gap: 5px; margin-top: 10px; justify-content: center;">
                    {dots_in_group}
                </div>
            </div>
        """

    html = f"""
    <div style="padding: 20px; background: #f0f2f6; border-radius: 15px; margin: 10px 0;">
        <h3 style="text-align: center; color: #333;">Dividing {dividend} into {divisor} Groups</h3>
        <p style="text-align: center; color: #555;">We are sharing {dividend} items equally among {divisor} groups.</p>
        <div style="display: grid; grid-template-columns: repeat(auto-fit, minmax(120px, 1fr)); gap: 15px; margin-top: 20px;">
            {groups_html}
        </div>
        <h4 style="text-align: center; margin-top: 25px; color: #333;">
            Each group gets <b>{quotient}</b> items. So, {dividend} ÷ {divisor} = {quotient}.
        </h4>
    </div>
    """
    return html

def create_counting_blocks(num1, num2, operation):
    """(Dynamic) Create colorful counting blocks for addition/subtraction."""
    html = f"""
    <div style="padding: 20px; background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); border-radius: 15px; margin: 10px 0;">
        <h3 style="color: white; text-align: center; margin-bottom: 20px;">🧮 Counting Blocks: {num1} {'+' if operation == 'add' else '−'} {num2}</h3>
        <div style="display: flex; justify-content: center; align-items: center; gap: 20px; flex-wrap: wrap;">
            <!-- Blocks for Num1 -->
            <div style="display: flex; flex-wrap: wrap; gap: 5px; border: 2px dashed #FFE066; padding: 5px; border-radius: 5px; align-items: center; justify-content: center; min-width: 100px;"><div style="width: 100%; text-align:center; color: white; font-weight: bold;">{num1}</div>{''.join([f'<div style="width: 25px; height: 25px; background: #FF6B6B; border-radius: 5px;"></div>' for _ in range(num1)])}</div>
            <div style="font-size: 40px; color: #FFE066;">{'+' if operation == 'add' else '−'}</div>
            <!-- Blocks for Num2 -->
            <div style="display: flex; flex-wrap: wrap; gap: 5px; border: 2px dashed #FFE066; padding: 5px; border-radius: 5px; align-items: center; justify-content: center; min-width: 100px;"><div style="width: 100%; text-align:center; color: white; font-weight: bold;">{num2}</div>{''.join([f'<div style="width: 25px; height: 25px; background: #4ECDC4; border-radius: 5px;"></div>' for _ in range(num2)])}</div>
            <div style="font-size: 40px; color: #FFE066;">=</div>
            <!-- Blocks for Answer -->
            <div style="display: flex; flex-wrap: wrap; gap: 5px; border: 2px solid white; background: rgba(255,255,255,0.2); padding: 5px; border-radius: 5px; align-items: center; justify-content: center; min-width: 100px;"><div style="width: 100%; text-align:center; color: white; font-weight: bold;">{num1 + num2 if operation == 'add' else max(0, num1 - num2)}</div>{''.join([f'<div style="width: 25px; height: 25px; background: #95E1D3; border-radius: 5px;"></div>' for _ in range(num1 + num2 if operation == 'add' else max(0, num1 - num2))])}</div>
        </div>
    </div>"""
    return html

def create_dynamic_fraction_circle(numerator, denominator):
    """(Dynamic) Generates an SVG of a pizza/pie to represent a fraction."""
    if not (0 < numerator <= denominator): return "<p>I can only show proper fractions!</p>"
    width, height, radius = 150, 150, 60
    cx, cy = width / 2, height / 2
    slices_html = ''
    angle_step = 360 / denominator
    for i in range(denominator):
        start_angle, end_angle = i * angle_step, (i + 1) * angle_step
        fill_color = "#FF6B6B" if i < numerator else "#DDDDDD"
        start_rad, end_rad = math.radians(start_angle - 90), math.radians(end_angle - 90)
        x1, y1 = cx + radius * math.cos(start_rad), cy + radius * math.sin(start_rad)
        x2, y2 = cx + radius * math.cos(end_rad), cy + radius * math.sin(end_rad)
        large_arc_flag = 1 if angle_step > 180 else 0
        path_d = f"M {cx},{cy} L {x1},{y1} A {radius},{radius} 0 {large_arc_flag},1 {x2},{y2} Z"
        slices_html += f'<path d="{path_d}" fill="{fill_color}" stroke="#333" stroke-width="2"/>'
    html = f"""<div style="padding: 20px; background: linear-gradient(135deg, #A8EDEA 0%, #FED6E3 100%); border-radius: 15px; margin: 10px 0;"><h3 style="color: #333; text-align: center;">Fraction Pizza: {numerator}/{denominator}</h3><div style="display: flex; justify-content: center;"><svg width="{width}" height="{height}">{slices_html}</svg></div><p style="color: #333; text-align: center; margin-top: 15px; font-size: 18px;">The pizza is cut into <b>{denominator}</b> equal slices, and we are showing <b>{numerator}</b> of them! 🍕</p></div>"""
    return html

def create_clock_visual(hours, minutes):
    """(Dynamic) Create a clock showing a specific time."""
    min_angle = minutes * 6
    hour_angle = (hours % 12 + minutes / 60) * 30
    html = f"""<div style="padding: 20px; background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); border-radius: 15px; margin: 10px 0;"><h3 style="color: white; text-align: center; margin-bottom: 20px;">🕐 Learning Time!</h3><div style="display: flex; justify-content: center;"><svg width="250" height="250" viewBox="0 0 250 250" style="background: white; border-radius: 50%; border: 8px solid #FFE066;"><circle cx="125" cy="125" r="110" fill="white" stroke="#333" stroke-width="2"/><text x="125" y="45" text-anchor="middle" font-size="20" font-weight="bold" fill="#333">12</text><text x="205" y="130" text-anchor="middle" font-size="20" font-weight="bold" fill="#333">3</text><text x="125" y="215" text-anchor="middle" font-size="20" font-weight="bold" fill="#333">6</text><text x="45" y="130" text-anchor="middle" font-size="20" font-weight="bold" fill="#333">9</text><line x1="125" y1="125" x2="125" y2="40" stroke="#FF6B6B" stroke-width="6" stroke-linecap="round" transform="rotate({hour_angle}, 125, 125)"/><line x1="125" y1="125" x2="125" y2="25" stroke="#4ECDC4" stroke-width="4" stroke-linecap="round" transform="rotate({min_angle}, 125, 125)"/><circle cx="125" cy="125" r="8" fill="#333"/></svg></div><div style="text-align: center; margin-top: 20px;"><p style="color: #FFE066; font-size: 24px; font-weight: bold;">This clock shows {hours:02d}:{minutes:02d}</p><p style="color: white; font-size: 16px;">The short <span style="color:#FF6B6B">red</span> hand points to the hour. The long <span style="color:#4ECDC4">blue</span> hand points to the minutes.</p></div></div>"""
    return html

def create_multiplication_array(rows, cols):
    """(Dynamic) Generates an SVG grid of dots to show small multiplication."""
    cell_size, gap = 25, 5
    svg_width = cols * (cell_size + gap)
    svg_height = rows * (cell_size + gap)
    dots_html = "".join([f'<circle cx="{c * (cell_size + gap) + cell_size/2}" cy="{r * (cell_size + gap) + cell_size/2}" r="{cell_size/2 - 2}" fill="#FF6B6B"/>' for r in range(rows) for c in range(cols)])
    html = f"""<div style="padding: 20px; background: linear-gradient(135deg, #FF9A9E 0%, #FECFEF 100%); border-radius: 15px; margin: 10px 0;"><h3 style="color:#333; text-align: center;">Multiplication Array: {rows} × {cols} = {rows * cols}</h3><div style="display: flex; justify-content: center; padding: 10px;"><svg width="{svg_width}" height="{svg_height}">{dots_html}</svg></div><p style="color: #333; text-align: center; font-size: 18px;">See? There are <b>{rows}</b> rows of <b>{cols}</b> dots. That's <b>{rows*cols}</b> dots in total!</p></div>"""
    return html

def create_number_line(start, end, points, title="Number Line"):
    """(Dynamic) Creates a simple number line SVG."""
    width = 600
    padding = 30
    if start >= end: end = start + 1
    scale = (width - 2 * padding) / (end - start)
    def to_x(n): return padding + (n - start) * scale
    ticks_html = "".join([f'<g transform="translate({to_x(i)}, 50)"><line y2="10" stroke="#aaa"/><text y="30" text-anchor="middle" fill="#555">{i}</text></g>' for i in range(start, end + 1)])
    points_html = "".join([f'<g transform="translate({to_x(p)}, 50)"><circle r="8" fill="#FF6B6B" stroke="white" stroke-width="2"/><text y="-15" text-anchor="middle" font-weight="bold" fill="#D63031">{p}</text></g>' for p in points])
    html = f"""<div style="padding: 20px; background: #f7f1e3; border-radius: 15px; margin: 10px 0;"><h3 style="text-align: center; color: #333;">{title}</h3><svg width="{width}" height="100"><line x1="{padding}" y1="50" x2="{width-padding}" y2="50" stroke="#333" stroke-width="2"/>{ticks_html}{points_html}</svg></div>"""
    return html

def create_place_value_blocks(number):
    """(FIXED & Dynamic) Create place value blocks for understanding numbers."""
    hundreds, tens, ones = number // 100, (number % 100) // 10, number % 10
    h_block_html, t_block_html, o_block_html = "", "", ""
    if hundreds > 0:
        hundreds_grid = "".join(["<div style='background:#F5A6A6'></div>"] * 100)
        hundreds_squares = "".join([f'<div style="width: 100px; height: 100px; background: #FF6B6B; border: 2px solid #D63031; display: grid; grid-template-columns: repeat(10, 1fr); gap: 2px; padding: 2px;">{hundreds_grid}</div>' for _ in range(hundreds)])
        h_block_html = f'<div style="text-align: center;"><h4>Hundreds: {hundreds}</h4><div style="display: flex; gap: 5px;">{hundreds_squares}</div></div>'
    if tens > 0:
        tens_grid = "".join(["<div style='background:#A2E8E4'></div>"] * 10)
        tens_sticks = "".join([f'<div style="width: 10px; height: 100px; background: #4ECDC4; border: 2px solid #00B894; display: grid; grid-template-rows: repeat(10, 1fr); gap: 2px; padding: 2px;">{tens_grid}</div>' for _ in range(tens)])
        t_block_html = f'<div style="text-align: center;"><h4>Tens: {tens}</h4><div style="display: flex; gap: 5px; align-items: flex-end;">{tens_sticks}</div></div>'
    if ones > 0:
         ones_cubes = "".join(['<div style="width: 10px; height: 10px; background: #FFE066; border: 2px solid #FDCB6E;"></div>' for _ in range(ones)])
         o_block_html = f'<div style="text-align: center;"><h4>Ones: {ones}</h4><div style="display: flex; gap: 5px; align-items: flex-end; flex-wrap: wrap; width: 50px; justify-content: center;">{ones_cubes}</div></div>'
    html = f"""<div style="padding: 20px; background: linear-gradient(135deg, #dfe6e9 0%, #b2bec3 100%); border-radius: 15px; margin: 10px 0;"><h3 style="color: #333; text-align: center;">Place Value Blocks for {number}</h3><div style="display: flex; justify-content: center; align-items: flex-end; gap: 20px; flex-wrap: wrap; padding: 20px 0; min-height: 150px;">{h_block_html}{t_block_html}{o_block_html}</div><div style="text-align: center; margin-top: 15px; padding: 10px; background: rgba(0,0,0,0.1); border-radius: 10px;"><h4 style="color: #333; margin:0;">{hundreds} Hundreds + {tens} Tens + {ones} Ones = {number}</h4></div></div>"""
    return html

def create_shape_explorer():
    """(Static) Create colorful shape recognition tool."""
    html = """<div style="padding: 20px; background: linear-gradient(135deg, #A8EDEA 0%, #FED6E3 100%); border-radius: 15px; margin: 10px 0;"><h3 style="color: #333; text-align: center; margin-bottom: 20px;">🔷 Shape Explorer!</h3><div style="display: grid; grid-template-columns: repeat(auto-fit, minmax(150px, 1fr)); gap: 20px; max-width: 600px; margin: 0 auto;"><div style="text-align: center; padding: 15px; background: white; border-radius: 10px; box-shadow: 0 4px 8px rgba(0,0,0,0.1);"><h4 style="color: #333; margin-bottom: 10px;">Circle</h4><svg width="80" height="80"><circle cx="40" cy="40" r="35" fill="#FF6B6B" stroke="#333" stroke-width="3"/></svg><p style="color: #666; font-size: 12px; margin-top: 10px;">Round and smooth!</p></div><div style="text-align: center; padding: 15px; background: white; border-radius: 10px; box-shadow: 0 4px 8px rgba(0,0,0,0.1);"><h4 style="color: #333; margin-bottom: 10px;">Square</h4><svg width="80" height="80"><rect x="12.5" y="12.5" width="55" height="55" fill="#4ECDC4" stroke="#333" stroke-width="3"/></svg><p style="color: #666; font-size: 12px; margin-top: 10px;">4 equal sides!</p></div><div style="text-align: center; padding: 15px; background: white; border-radius: 10px; box-shadow: 0 4px 8px rgba(0,0,0,0.1);"><h4 style="color: #333; margin-bottom: 10px;">Triangle</h4><svg width="80" height="80"><polygon points="40,15 15,65 65,65" fill="#FFD93D" stroke="#333" stroke-width="3"/></svg><p style="color: #666; font-size: 12px; margin-top: 10px;">3 sides and corners!</p></div><div style="text-align: center; padding: 15px; background: white; border-radius: 10px; box-shadow: 0 4px 8px rgba(0,0,0,0.1);"><h4 style="color: #333; margin-bottom: 10px;">Rectangle</h4><svg width="80" height="80"><rect x="10" y="25" width="60" height="30" fill="#95E1D3" stroke="#333" stroke-width="3"/></svg><p style="color: #666; font-size: 12px; margin-top: 10px;">4 sides, opposite sides equal!</p></div></div><p style="color: #333; text-align: center; margin-top: 20px; font-size: 18px;">Can you find these shapes around you? 🔍✨</p></div>"""
    return html

def create_money_counter():
    """(Static) Create coin counting visual."""
    html = """<div style="padding: 20px; background: linear-gradient(135deg, #FFE259 0%, #FFA751 100%); border-radius: 15px; margin: 10px 0;"><h3 style="color: #333; text-align: center; margin-bottom: 20px;">💰 Money Counter!</h3><div style="display: flex; justify-content: center; gap: 30px; flex-wrap: wrap;"><div style="text-align: center; padding: 15px; background: white; border-radius: 10px;"><h4 style="color: #333;">Penny</h4><div style="width: 50px; height: 50px; background: #CD7F32; border-radius: 50%; margin: 10px auto; display: flex; align-items: center; justify-content: center; border: 3px solid #8B4513;"><span style="color: white; font-weight: bold;">1¢</span></div><p style="color: #666; font-size: 12px;">1 cent</p></div><div style="text-align: center; padding: 15px; background: white; border-radius: 10px;"><h4 style="color: #333;">Nickel</h4><div style="width: 55px; height: 55px; background: #C0C0C0; border-radius: 50%; margin: 10px auto; display: flex; align-items: center; justify-content: center; border: 3px solid #808080;"><span style="color: #333; font-weight: bold;">5¢</span></div><p style="color: #666; font-size: 12px;">5 cents</p></div><div style="text-align: center; padding: 15px; background: white; border-radius: 10px;"><h4 style="color: #333;">Dime</h4><div style="width: 45px; height: 45px; background: #C0C0C0; border-radius: 50%; margin: 10px auto; display: flex; align-items: center; justify-content: center; border: 3px solid #808080;"><span style="color: #333; font-weight: bold;">10¢</span></div><p style="color: #666; font-size: 12px;">10 cents</p></div><div style="text-align: center; padding: 15px; background: white; border-radius: 10px;"><h4 style="color: #333;">Quarter</h4><div style="width: 60px; height: 60px; background: #C0C0C0; border-radius: 50%; margin: 10px auto; display: flex; align-items: center; justify-content: center; border: 3px solid #808080;"><span style="color: #333; font-weight: bold;">25¢</span></div><p style="color: #666; font-size: 12px;">25 cents</p></div></div><p style="color: #333; text-align: center; margin-top: 20px; font-size: 18px;">Practice counting coins to make different amounts! 🪙✨</p></div>"""
    return html

# --- [The rest of your application code remains the same] ---
# --- API KEY & MODEL CONFIGURATION, SESSION STATE, DIALOGS, etc. ---

# --- API KEY & MODEL CONFIGURATION ---
load_dotenv()
api_key = None

try:
    api_key = st.secrets["GOOGLE_API_KEY"]
except (KeyError, FileNotFoundError):
    api_key = os.getenv("GOOGLE_API_KEY")

if api_key:
    genai.configure(api_key=api_key)
    
    # Main text model
    model = genai.GenerativeModel(
        model_name="gemini-1.5-flash",
        system_instruction="""
        You are "Math Jegna", an AI specializing exclusively in K-12 mathematics.
        Your one and only function is to solve and explain math problems for children.
        You are an AI math tutor that uses the Professor B methodology developed by Everard Barrett. This methodology is designed to activate children's natural learning capacities and present mathematics as a contextual, developmental story that makes sense.

        IMPORTANT: When explaining mathematical concepts to young learners, mention that colorful visual aids will be provided to help illustrate the concept. Use phrases like:
        - "Let me show you this with some colorful blocks..."
        - "A fun visual will help you see how this works..."
        - "Let's use an area model to understand this multiplication problem..."
        - "I'll create a picture showing how we can divide these into groups..."
        
        Focus on concepts appropriate for K-12 students:
        - Basic counting and number recognition
        - Simple addition and subtraction (using manipulatives)
        - Multiplication as arrays or groups
        - Division as sharing into equal groups
        - Basic shapes and geometry
        - Place value with hundreds, tens, ones
        - Money counting and coin recognition
        - Time telling with analog clocks
        
        Always use age-appropriate language and relate math to real-world examples children understand.
        
Core Philosophy and Principles
1. Contextual Learning Approach
Present math as a story: Every mathematical concept should be taught as part of a continuing narrative that builds connections between ideas
Use concrete manipulatives: Always relate abstract concepts to physical, visual representations
Truth-telling: Present arithmetic computations simply and truthfully without confusing steps

2. Natural Learning Activation
Leverage natural capacities: Recognize that each child has mental capabilities designed to learn naturally
Story-based retention: Use stories and visual representations that children can easily remember
Reduced anxiety: Make math fun and engaging, not scary or confusing

3. Hands-on Learning
Mental gymnastics: Use finger counting, visual blocks, and interactive elements
No rote memorization: Focus on understanding through play and exploration
Build confidence: Celebrate small victories and progress

        You are strictly forbidden from answering any question that is not mathematical in nature.
        If you receive a non-mathematical question, you MUST decline with: "I can only answer math questions for students. Please ask me about numbers, shapes, counting, or other math topics!"
        
        Keep explanations simple, encouraging, and fun for young learners.
        """
    )
else:
    st.error("🚨 Google API Key not found! Please add it to your secrets or a local .env file.")
    st.stop()

# --- SESSION STATE & LOCAL STORAGE INITIALIZATION ---
if "chats" not in st.session_state:
    try:
        shared_chat_b64 = st.query_params.get("shared_chat")
        if shared_chat_b64:
            decoded_chat_json = base64.urlsafe_b64decode(shared_chat_b64).decode()
            st.session_state.chats = {"Shared Chat": json.loads(decoded_chat_json)}
            st.session_state.active_chat_key = "Shared Chat"
            st.query_params.clear()
        else:
            raise ValueError("No shared chat")
    except (TypeError, ValueError, Exception):
        saved_data_json = localS.getItem("math_mentor_chats")
        if saved_data_json:
            saved_data = json.loads(saved_data_json)
            st.session_state.chats = saved_data.get("chats", {})
            st.session_state.active_chat_key = saved_data.get("active_chat_key", "New Chat")
        else:
            st.session_state.chats = {
                "New Chat": [
                    {"role": "assistant", "content": "Hello! I'm Math Jegna, your friendly math helper! 🧠✨ I love helping students learn math with colorful pictures and fun activities. What would you like to learn about today? Maybe counting, shapes, or solving a math problem? 🌟"}
                ]
            }
            st.session_state.active_chat_key = "New Chat"

# --- RENAME DIALOG ---
@st.dialog("Rename Chat")
def rename_chat(chat_key):
    st.write(f"Enter a new name for '{chat_key}':")
    new_name = st.text_input("New Name", key=f"rename_input_{chat_key}")
    if st.button("Save", key=f"save_rename_{chat_key}"):
        if new_name and new_name not in st.session_state.chats:
            st.session_state.chats[new_name] = st.session_state.chats.pop(chat_key)
            st.session_state.active_chat_key = new_name
            st.rerun()
        elif not new_name:
            st.error("Name cannot be empty.")
        else:
            st.error("A chat with this name already exists.")

# --- DELETE CONFIRMATION DIALOG ---
@st.dialog("Delete Chat")
def delete_chat(chat_key):
    st.warning(f"Are you sure you want to delete '{chat_key}'? This cannot be undone.")
    if st.button("Yes, Delete", type="primary", key=f"confirm_delete_{chat_key}"):
        st.session_state.chats.pop(chat_key)
        # Add the logic to switch to a new or different chat after deletion
        if st.session_state.active_chat_key == chat_key:
            # Simple fallback to the first available chat or a new one
            if st.session_state.chats:
                st.session_state.active_chat_key = next(iter(st.session_state.chats))
            else:
                # Create a new chat if none are left
                st.session_state.chats["New Chat"] = [
                    {"role": "assistant", "content": "Hello! Let's start a new math adventure! 🚀"}
                ]
                st.session_state.active_chat_key = "New Chat"
        st.rerun()

# --- MAIN APP LAYOUT ---
with st.sidebar:
    st.title("🧮 Math Jegna")
    st.write("Your K-8 AI Math Tutor")
    st.divider()
    
    # Chat history list
    for chat_key in list(st.session_state.chats.keys()):
        col1, col2, col3 = st.columns([0.6, 0.2, 0.2])
        with col1:
            if st.button(chat_key, key=f"switch_{chat_key}", use_container_width=True, type="primary" if st.session_state.active_chat_key == chat_key else "secondary"):
                st.session_state.active_chat_key = chat_key
                st.rerun()
        with col2:
            if st.button("✏️", key=f"rename_{chat_key}", help="Rename Chat"):
                rename_chat(chat_key)
        with col3:
            if st.button("🗑️", key=f"delete_{chat_key}", help="Delete Chat"):
                 delete_chat(chat_key)

    if st.button("➕ New Chat", use_container_width=True):
        new_chat_name = f"Chat {len(st.session_state.chats) + 1}"
        # Ensure the name is unique
        while new_chat_name in st.session_state.chats:
            new_chat_name += "*"
        st.session_state.chats[new_chat_name] = [
             {"role": "assistant", "content": "Ready for a new math problem! What's on your mind? 😃"}
        ]
        st.session_state.active_chat_key = new_chat_name
        st.rerun()
        
    st.divider()
    
    # Save chats to local storage
    if st.button("💾 Save Chats", use_container_width=True):
        data_to_save = {
            "chats": st.session_state.chats,
            "active_chat_key": st.session_state.active_chat_key
        }
        localS.setItem("math_mentor_chats", json.dumps(data_to_save))
        st.toast("Chats saved to your browser!", icon="✅")

    # Download chat button
    active_chat_history = st.session_state.chats[st.session_state.active_chat_key]
    download_str = format_chat_for_download(active_chat_history)
    st.download_button(
        label="📥 Download Chat",
        data=download_str,
        file_name=f"{st.session_state.active_chat_key.replace(' ', '_')}_history.md",
        mime="text/markdown",
        use_container_width=True
    )
    
    # Share chat button
    if st.button("🔗 Share Chat", use_container_width=True):
        chat_json = json.dumps(st.session_state.chats[st.session_state.active_chat_key])
        chat_b64 = base64.urlsafe_b64encode(chat_json.encode()).decode()
        share_url = f"https://huggingface.co/spaces/YOUR_SPACE_HERE?shared_chat={chat_b64}" # Placeholder
        st.code(share_url)
        st.info("Copy the URL above to share this specific chat! (You might need to update the base URL)")


st.header(f"Chatting with Math Jegna: _{st.session_state.active_chat_key}_")

# Display chat messages
for message in st.session_state.chats[st.session_state.active_chat_key]:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])
        # If a visual was generated and saved with the message, display it
        if "visual_html" in message and message["visual_html"]:
            components.html(message["visual_html"], height=450, scrolling=True)

# User input
if prompt := st.chat_input("Ask a K-8 math question..."):
    # Add user message to chat history
    st.session_state.chats[st.session_state.active_chat_key].append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    # Prepare chat for Gemini API
    gemini_chat_history = [
        {"role": convert_role_for_gemini(m["role"]), "parts": [m["content"]]}
        for m in st.session_state.chats[st.session_state.active_chat_key]
    ]

    # Generate response
    with st.chat_message("assistant"):
        with st.spinner("Math Jegna is thinking..."):
            try:
                chat_session = model.start_chat(history=gemini_chat_history)
                response = chat_session.send_message(prompt, stream=True)
                
                full_response = ""
                response_container = st.empty()
                for chunk in response:
                    full_response += chunk.text
                    response_container.markdown(full_response + " ▌")
                response_container.markdown(full_response)
                
                # After generating text, decide if a visual is needed and generate it
                visual_html_content = None
                if should_generate_visual(prompt, full_response):
                    visual_html_content = create_visual_manipulative(prompt, full_response)
                    if visual_html_content:
                        components.html(visual_html_content, height=450, scrolling=True)

                # Add AI response and visual to session state
                st.session_state.chats[st.session_state.active_chat_key].append({
                    "role": "assistant", 
                    "content": full_response,
                    "visual_html": visual_html_content # Store the visual with the message
                })


            except genai.types.generation_types.BlockedPromptException as e:
                error_message = "I can only answer math questions for students. Please ask me about numbers, shapes, or other math topics!"
                st.error(error_message)
                st.session_state.chats[st.session_state.active_chat_key].append({"role": "assistant", "content": error_message, "visual_html": None})
            except Exception as e:
                st.error(f"An error occurred: {e}")