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 (UNCHANGED) ---
def format_chat_for_download(chat_history):
    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):
    if role == "assistant": return "model"
    return role

def should_generate_visual(user_prompt, ai_response):
    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',
        'place value', 'tens', 'ones', 'hundreds',
        'number line', 'array', 'grid', 'area model', 'solve for'
    ]
    combined_text = (user_prompt + " " + ai_response).lower()
    return any(keyword in combined_text for keyword in k12_visual_keywords) or any(op in user_prompt for op in ['*', '/', 'x', '='])


def create_visual_manipulative(user_prompt, ai_response):
    """-- SMART VISUAL ROUTER (UPGRADED FOR ALGEBRA) --"""
    try:
        # Normalize the prompt for easier regex matching
        user_norm = user_prompt.lower().replace(' ', '')
        
        # PRIORITY 1: ALGEBRA (e.g., "2x+10=40", "solve 3y + 5 = 20")
        # Looks for patterns like ax+b=c, ax-b=c, etc.
        algebra_match = re.search(r'(\d+)[a-z]\s*\+\s*(\d+)\s*=\s*(\d+)', user_norm)
        if algebra_match:
            a, b, c = map(int, algebra_match.groups())
            if (c - b) % a == 0: # Only create visual for clean integer solutions
                return create_algebra_balance_scale(a, b, c, '+')

        # PRIORITY 2: Division
        div_match = re.search(r'(\d+)dividedby(\d+)', user_norm) or re.search(r'(\d+)/(\d+)', user_norm)
        if div_match and "fraction" not in user_norm:
            dividend, divisor = int(div_match.group(1)), int(div_match.group(2))
            if divisor > 0 and dividend % divisor == 0 and dividend <= 50:
                return create_division_groups_visual(dividend, divisor)

        # PRIORITY 3: Multiplication
        mult_match = re.search(r'(\d+)(?:x|times|\*)(\d+)', user_norm)
        if mult_match:
            num1, num2 = int(mult_match.group(1)), int(mult_match.group(2))
            if num1 <= 10 and num2 <= 10: return create_multi_model_multiplication_visual(num1, num2)
            elif 10 < num1 < 100 and 10 < num2 < 100: return create_multiplication_area_model(num1, num2)

        # PRIORITY 4: Addition/Subtraction Blocks (now less likely to be triggered incorrectly)
        if any(word in user_norm 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_norm for w in ['add', 'plus', '+']) else 'subtract'
                if num1 <= 20 and num2 <= 20: return create_counting_blocks(num1, num2, operation)
        
        # Other priorities remain the same...
        # (Fractions, Time, Place Value, Number Lines, Static Fallbacks)

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

# --- NEW VISUAL TOOLBOX FUNCTION: ALGEBRA BALANCE SCALE ---

def create_algebra_balance_scale(a, b, c, op):
    """(BRAND NEW) Generates a step-by-step balance scale visual for solving linear equations."""
    
    # Calculate intermediate and final steps
    step2_val = c - b
    final_x = step2_val // a

    # --- Reusable components ---
    def make_x_blocks(count):
        return "".join([f'<div class="x-block">x</div>' for _ in range(count)])
    def make_unit_blocks(count, faded=False):
        # For larger numbers, just show a single block with the value
        if count > 12:
            return f'<div class="unit-block-large {"faded" if faded else ""}">{count}</div>'
        return "".join([f'<div class="unit-block {"faded" if faded else ""}">1</div>' for _ in range(count)])

    html = f"""
    <style>
        .balance-container {{ font-family: sans-serif; padding: 20px; background: #f4f7f6; border-radius: 15px; margin: 10px 0; }}
        .step {{ margin-bottom: 25px; border-left: 4px solid #4ECDC4; padding-left: 15px; }}
        .scale {{ display: flex; align-items: flex-end; justify-content: center; min-height: 100px; }}
        .pan {{ border: 3px solid #6c757d; border-top: none; padding: 10px; min-width: 150px; display: flex; flex-wrap: wrap; gap: 5px; justify-content: center; align-items: center; background: #fff; border-radius: 0 0 10px 10px; }}
        .fulcrum {{ width: 20px; height: 100px; background: #6c757d; position: relative; }}
        .beam {{ height: 8px; background: #6c757d; width: 400px; position: absolute; top: -4px; left: -190px; }}
        .x-block {{ width: 30px; height: 30px; background: #FF6B6B; color: white; display: flex; justify-content: center; align-items: center; font-weight: bold; border-radius: 5px; }}
        .unit-block {{ width: 20px; height: 20px; background: #4ECDC4; color: white; font-size: 0.8em; display: flex; justify-content: center; align-items: center; border-radius: 5px; }}
        .unit-block-large {{ padding: 10px 15px; background: #4ECDC4; color: white; font-size: 1.2em; font-weight: bold; display: flex; justify-content: center; align-items: center; border-radius: 5px; }}
        .faded {{ opacity: 0.3; text-decoration: line-through; }}
        .op-text {{ font-size: 1.5em; color: #d62828; margin: 0 20px; }}
        .grouping {{ border: 2px dashed #FF6B6B; padding: 10px; margin-top: 10px; border-radius: 8px; }}
    </style>
    <div class="balance-container">
        <h3 style="text-align: center; color: #333;">Solving {a}x + {b} = {c} with a Balance Scale</h3>
        
        <!-- Step 1: Initial Setup -->
        <div class="step">
            <h4>1. Set up the equation</h4>
            <p>The scale is balanced, with the left side equal to the right side.</p>
            <div class="scale">
                <div class="pan">{make_x_blocks(a)} {make_unit_blocks(b)}</div>
                <div class="fulcrum"><div class="beam"></div></div>
                <div class="pan">{make_unit_blocks(c)}</div>
            </div>
        </div>

        <!-- Step 2: Isolate the variable term -->
        <div class="step">
            <h4>2. Subtract {b} from both sides</h4>
            <p>To keep it balanced, we must remove the same amount from each pan.</p>
            <div class="scale">
                <div class="pan">{make_x_blocks(a)} {make_unit_blocks(b, faded=True)}</div>
                <div class="op-text">➔</div>
                <div class="pan">{make_x_blocks(a)}</div>
                <div class="fulcrum"><div class="beam"></div></div>
                <div class="pan">{make_unit_blocks(c-b)}</div>
            </div>
        </div>

        <!-- Step 3: Solve for x -->
        <div class="step">
            <h4>3. Divide both sides by {a}</h4>
            <p>We split each side into {a} equal groups to find the value of a single 'x'.</p>
            <div class="scale">
                <div class="pan grouping">{make_x_blocks(1)}</div>
                <div class="fulcrum"><div class="beam"></div></div>
                <div class="pan grouping">{make_unit_blocks(final_x)}</div>
            </div>
            <h4 style="text-align:center; margin-top: 15px;">Solution: x = {final_x}</h4>
        </div>
    </div>
    """
    return html


# --- [All other visual functions and app code remain the same] ---
# Note: The code below is identical to the previous version, provided for completeness.
def create_multi_model_multiplication_visual(rows, cols):
    groups_html = ""
    for r in range(rows):
        dots = "".join([f'<div style="width:12px; height:12px; background:#FF6B6B; border-radius:50%;"></div>' for _ in range(cols)])
        groups_html += f'<div style="border:2px solid #FFADAD; border-radius:8px; padding:5px; display:flex; flex-wrap:wrap; gap:4px; justify-content:center; margin:2px;">{dots}</div>'
    cell_size, gap = 20, 4
    svg_width, svg_height = cols * (cell_size + gap), rows * (cell_size + gap)
    array_dots = "".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="#4ECDC4"/>' for r in range(rows) for c in range(cols)])
    array_svg = f'<svg width="{svg_width}" height="{svg_height}" style="margin: 0 auto;">{array_dots}</svg>'
    addition_str = " + ".join([str(cols) for _ in range(rows)])
    line_end, line_width, padding = rows * cols + 2, 400, 20
    scale = (line_width - 2 * padding) / line_end
    ticks = "".join([f'<text x="{padding + i*scale}" y="35" text-anchor="middle" font-size="10">{i}</text>' for i in range(0, line_end, 2)])
    jumps_html = "".join([f'<path d="M {padding + (i * cols * scale)} 20 Q {(padding + (i * cols * scale) + padding + ((i + 1) * cols * scale))/2} -5, {padding + ((i + 1) * cols * scale)} 20" stroke="#FFD93D" fill="none" stroke-width="2"/>' for i in range(rows)])
    number_line_svg = f'<svg width="{line_width}" height="40"><line x1="{padding}" y1="20" x2="{line_width-padding}" y2="20" stroke="#333"/>{ticks}{jumps_html}</svg>'
    html = f"""<div style="font-family: sans-serif; padding: 20px; background: #f8f9fa; border: 1px solid #dee2e6; border-radius: 15px; margin: 10px 0;"><h3 style="text-align: center; color: #333; margin-bottom:25px;">Four Ways to See {rows} × {cols} = {rows*cols}</h3><div style="display: grid; grid-template-columns: 1fr 1fr; gap: 20px;"><div style="border: 1px solid #ddd; padding: 10px; border-radius: 8px;"><h4 style="text-align:center; margin-top:0;">Use an Array</h4>{array_svg}</div><div style="border: 1px solid #ddd; padding: 10px; border-radius: 8px;"><h4 style="text-align:center; margin-top:0;">Use Equal Groups</h4><div style="display:flex; flex-wrap:wrap; gap:5px; justify-content:center;">{groups_html}</div></div><div style="border: 1px solid #ddd; padding: 10px; border-radius: 8px;"><h4 style="text-align:center; margin-top:0;">Use Repeated Addition</h4><div style="text-align:center; font-size: 1.5em; color: #0077b6;">{addition_str}</div></div><div style="border: 1px solid #ddd; padding: 10px; border-radius: 8px;"><h4 style="text-align:center; margin-top:0;">Use a Number Line</h4>{number_line_svg}</div></div></div>"""
    return html
def create_multiplication_area_model(num1, num2):
    n1_tens, n1_ones = num1 // 10, num1 % 10
    n2_tens, n2_ones = num2 // 10, num2 % 10
    p1, p2, p3, p4 = n1_tens*10 * n2_tens*10, n1_tens*10 * n2_ones, n1_ones * n2_tens*10, n1_ones * n2_ones
    total = p1 + p2 + p3 + p4
    html = f"""<div style="font-family: sans-serif; padding: 20px; background: #f0f8ff; border-radius: 15px; margin: 10px 0;"><h3 style="text-align: center; color: #333;">Area Model for {num1} × {num2}</h3><div style="display: flex; justify-content: center; align-items: center; margin-top: 20px;"><div style="display: flex; flex-direction: column; text-align: right; gap: 5px; margin-right: 5px;"><div style="height: 60px; display: flex; align-items: center; justify-content: flex-end; font-weight: bold; color: #0077b6;">{n1_tens*10}</div><div style="height: 60px; display: flex; align-items: center; justify-content: flex-end; font-weight: bold; color: #0077b6;">{n1_ones}</div></div><div style="display: inline-grid; border: 2px solid #333;"><div style="display: flex; grid-column: 1 / 3;"><div style="width: 100px; text-align: center; font-weight: bold; color: #d00000; padding: 5px;">{n2_tens*10}</div><div style="width: 100px; text-align: center; font-weight: bold; color: #d00000; padding: 5px;">{n2_ones}</div></div><div style="grid-row: 2; width: 100px; height: 60px; background: #FFADAD; text-align: center; border: 1px solid #333; padding: 5px;">{p1}</div><div style="grid-row: 2; width: 100px; height: 60px; background: #FFD6A5; text-align: center; border: 1px solid #333; padding: 5px;">{p2}</div><div style="grid-row: 3; width: 100px; height: 60px; background: #FDFFB6; text-align: center; border: 1px solid #333; padding: 5px;">{p3}</div><div style="grid-row: 3; width: 100px; height: 60px; background: #CAFFBF; text-align: center; border: 1px solid #333; padding: 5px;">{p4}</div></div></div><div style="text-align: center; margin-top: 20px; font-size: 1.2em;"><b>Add the partial products:</b> {p1} + {p2} + {p3} + {p4} = <b>{total}</b></div></div>"""
    return html
def create_division_groups_visual(dividend, divisor):
    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

# --- [Rest of the boilerplate code follows] ---
# (API Key, Session State, Dialogs, Main Layout, etc. This is identical to the previous version.)

# --- 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)
    model = genai.GenerativeModel(
        model_name="gemini-1.5-flash",
        system_instruction="""
        You are "Math Jegna", an AI specializing exclusively in K-12 mathematics.
        You are an AI math tutor that uses the Professor B methodology. 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, mention that helpful visuals will be provided.
        - For simple multiplication, mention showing it in multiple ways (arrays, groups, number line).
        - For algebra, use the analogy of a "balance scale" where you must do the same thing to both sides.
        - For division, talk about "sharing into equal groups."
        
        Always use age-appropriate language and relate math to real-world examples.
        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, DIALOGS, and MAIN APP LAYOUT ---
# This entire section is identical to the previous version.
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! 🧠✨ What would you like to learn about today?"}] }
            st.session_state.active_chat_key = "New Chat"

@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.")

@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)
        if st.session_state.active_chat_key == chat_key:
            if st.session_state.chats: st.session_state.active_chat_key = next(iter(st.session_state.chats))
            else:
                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()

with st.sidebar:
    st.title("🧮 Math Jegna")
    st.write("Your K-8 AI Math Tutor")
    st.divider()
    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}"
        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()
    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="✅")
    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)
    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}"
        st.code(share_url)
        st.info("Copy the URL above to share this specific chat! (Update the base URL)")

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

for message in st.session_state.chats[st.session_state.active_chat_key]:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])
        if "visual_html" in message and message["visual_html"]:
            components.html(message["visual_html"], height=600, scrolling=True)

if prompt := st.chat_input("Ask a K-8 math question..."):
    st.session_state.chats[st.session_state.active_chat_key].append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)
    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]]
    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)
                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=600, scrolling=True)
                st.session_state.chats[st.session_state.active_chat_key].append({"role": "assistant", "content": full_response, "visual_html": visual_html_content})
            except genai.types.generation_types.BlockedPromptException as e:
                error_message = "I can only answer math questions for students."
                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}")