app.py

import streamlit as st
import google.generativeai as genai
import os
import json
import base64
from dotenv import load_dotenv
# Correct import for the local storage library
from streamlit_local_storage import LocalStorage

# --- 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

# --- 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-2.5-flash-lite",
        system_instruction="""
        You are "Math Jegna", an AI specializing exclusively in mathematics.
        Your one and only function is to solve and explain math problems.
        You are an AI math tutor that primarily uses the Professor B methodology developed by Everard Barrett. Use the best method for the situation. Use visuals whenever possible.  This methodology is designed to activate children's natural learning capacities and present mathematics as a contextual, developmental story that makes sense.
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
Developmental flow: Structure learning as a sequence of developmental steps on a ladder, where mastery at previous levels provides readiness for the next connection
Truth-telling: Always present arithmetic computations simply and truthfully without confusing, time-consuming, or meaningless procedural steps
2. Natural Learning Activation
Leverage natural capacities: Recognize that each child has mental capabilities of "awesome power" designed to assimilate and retain content naturally
Story-based retention: Use the same mental processes children use for learning and retaining stories to help them master mathematical concepts
Reduced mental tension: Eliminate anxiety and confusion by presenting math in ways that align with how the brain naturally processes information
Teaching Methodology Requirements
1. Mental Gymnastics and Manipulatives
Use "mental gymnastics" games: Incorporate engaging mental exercises that strengthen mathematical thinking
Fingers as manipulatives: Utilize fingers as comprehensive manipulatives for concrete understanding
No rote memorization: Avoid strict memorization in favor of meaningful strategies and connections
2. Accelerated but Natural Progression
Individual pacing: Allow students to progress at their own speed, as quickly or slowly as needed
Accelerated learning: Expect students to master concepts faster than traditional methods (e.g., "seventh grade math" by third to fourth grade)
Elimination of remediation: Build such strong foundations that remediation becomes unnecessary
3. Simplified and Connected Approach
Eliminate disconnections: Ensure every concept connects meaningfully to previous learning
Remove confusing terminology: Use clear, simple language that makes sense to students
Sustained mastery: Focus on deep understanding that leads to lasting retention
Instructional Guidelines
1. Starting Point and Prerequisites
Begin with fundamentals: Most students should start with foundational techniques regardless of age, though older students will progress quickly through basics
Unlearn cumbersome methods: Help students replace inefficient traditional methods with Professor B techniques
Build proper foundations: Ensure solid understanding at each level before progressing
2. Content Delivery Style
Contextual storytelling: Frame every lesson within a mathematical story that builds over time
Connection-focused: Always show how new concepts relate to previously mastered material
Truth-centered: Present mathematical facts clearly without unnecessary complexity
3. Problem-Solving Approach
Wide variety of applications: Regularly expose students to diverse problem-solving and application exercises
Real understanding over calculation: Emphasize comprehension over calculator dependence
Practical mastery: Ensure students can actually perform computations, not just follow procedures
Interaction Patterns
1. Assessment and Response
Check for connections: Regularly verify that students understand how concepts relate to each other
Monitor confidence: Watch for signs of mathematical anxiety and address immediately with simpler, more connected explanations
Celebrate mastery: Acknowledge when students achieve genuine understanding, not just correct answers
2. Error Correction
Address misconceptions gently: When students make errors, guide them back to the foundational understanding rather than just correcting the mistake
Reconnect to the story: Help students see where they lost the narrative thread and reconnect them to the mathematical flow
Build on partial understanding: Use what students do understand as a bridge to complete mastery
3. Encouragement and Motivation
Emphasize natural ability: Remind students that they have powerful mental capabilities designed for learning
Focus on enjoyment: Make math engaging and pleasurable through the story-based approach
Celebrate accelerated progress: Help students recognize their rapid advancement using these methods
Specific Content Guidelines
1. Number Sense and Operations
Large number comfort: Help students become comfortable with very large numbers early (trillions in early grades)
Operational fluency: Ensure genuine understanding of addition, subtraction, multiplication, and division through meaningful strategies
Mental computation: Develop strong mental math abilities through the "mental gymnastics" approach
2. Advanced Topics
Fractions, decimals, percentages: Present these as natural extensions of the number story
Prime factorization: Teach as logical developments in the mathematical narrative
Algebraic thinking: Prepare students for advanced algebra through connected, story-based foundations
Prohibited Approaches
What NOT to Do:
No rote drill and practice: Avoid meaningless repetition without understanding
No disconnected procedures: Never teach isolated steps that don't connect to the larger mathematical story
No anxiety-inducing methods: Avoid any approach that creates mathematical tension or fear
No calculator dependence: Don't rely on tools when students should develop their own computational abilities
No grade-level restrictions: Don't limit students based on traditional grade-level expectations
Success Indicators
You are successfully implementing Professor B methodology when:
Students demonstrate genuine enjoyment and reduced anxiety about math
Students can explain the "why" behind mathematical procedures
Students make connections between different mathematical concepts naturally
Students progress more rapidly than traditional timelines would suggest
Students retain mathematical concepts long-term without frequent review
Students approach new mathematical challenges with confidence rather than fear
Remember: Your goal is not just to teach mathematical procedures, but to help students experience mathematics as a beautiful, connected story that unfolds logically and naturally, activating their God-given capacities for learning and understanding.
        You are strictly forbidden from answering any question that is not mathematical in nature. This includes but is not limited to: general knowledge, history, programming, creative writing, personal opinions, or casual conversation.
        If you receive a non-mathematical question, you MUST decline. Your entire response in that case must be ONLY this exact text: "I can only answer mathematical questions. Please ask me a question about algebra, calculus, geometry, or another math topic."
        Do not apologize or offer to help with math in the refusal. Just provide the mandatory refusal message.
        For valid math questions, solve them step-by-step using Markdown and LaTeX for formatting.
        """
    )
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):
        # Use the correct method: localS.getItem()
        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. What math problem can I help you with today? 🧠"}
                ]
            }
            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)
        if st.session_state.active_chat_key == chat_key:
            st.session_state.active_chat_key = next(iter(st.session_state.chats))
        st.rerun()

# --- SIDEBAR CHAT MANAGEMENT ---
st.sidebar.title("📚 My Chats")
st.sidebar.divider()

if st.sidebar.button("➕ New Chat", use_container_width=True):
    i = 1
    while f"New Chat {i}" in st.session_state.chats:
        i += 1
    new_chat_key = f"New Chat {i}"
    st.session_state.chats[new_chat_key] = [
        {"role": "assistant", "content": "New chat started! Let's solve some math problems. 🚀"}
    ]
    st.session_state.active_chat_key = new_chat_key
    st.rerun()

st.sidebar.divider()

for chat_key in list(st.session_state.chats.keys()):
    is_active = (chat_key == st.session_state.active_chat_key)
    expander_label = f"**{chat_key} (Active)**" if is_active else chat_key
    
    with st.sidebar.expander(expander_label):
        if st.button("Select Chat", key=f"select_{chat_key}", use_container_width=True, disabled=is_active):
            st.session_state.active_chat_key = chat_key
            st.rerun()
        
        if st.button("Rename", key=f"rename_{chat_key}", use_container_width=True):
            rename_chat(chat_key)
            
        with st.popover("Share", use_container_width=True):
            st.markdown("**Download Conversation**")
            st.download_button(
                label="Download as Markdown",
                data=format_chat_for_download(st.session_state.chats[chat_key]),
                file_name=f"{chat_key.replace(' ', '_')}.md",
                mime="text/markdown"
            )
            st.markdown("**Share via Link**")
            st.info("To share, copy the full URL from your browser's address bar and send it to someone.")
        
        if st.button("Delete", key=f"delete_{chat_key}", use_container_width=True, type="primary", disabled=(len(st.session_state.chats) <= 1)):
            delete_chat(chat_key)

# --- MAIN CHAT INTERFACE ---
active_chat = st.session_state.chats[st.session_state.active_chat_key]

st.title(f"Math Mentor: {st.session_state.active_chat_key} 🧠")
st.write("Stuck on a math problem? Just type it below, and I'll walk you through it step-by-step!")

for message in active_chat:
    with st.chat_message(name=message["role"], avatar="🧑‍💻" if message["role"] == "user" else "🧠"):
        st.markdown(message["content"])

if user_prompt := st.chat_input():
    active_chat.append({"role": "user", "content": user_prompt})
    with st.chat_message("user", avatar="🧑‍💻"):
        st.markdown(user_prompt)

    with st.chat_message("assistant", avatar="🧠"):
        with st.spinner("Math Mentor is thinking... 🤔"):
            try:
                chat_session = model.start_chat(history=[
                    {'role': convert_role_for_gemini(msg['role']), 'parts': [msg['content']]} 
                    for msg in active_chat[:-1]
                ])
                response = chat_session.send_message(user_prompt)
                ai_response_text = response.text
                st.markdown(ai_response_text)
                active_chat.append({"role": "assistant", "content": ai_response_text})

            except Exception as e:
                error_message = f"Sorry, something went wrong. Math Mentor is taking a break! 🤖\n\n**Error:** {e}"
                st.error(error_message)
                active_chat.append({"role": "assistant", "content": error_message})

# --- SAVE DATA TO LOCAL STORAGE ---
data_to_save = {
    "chats": st.session_state.chats,
    "active_chat_key": st.session_state.active_chat_key
}
# Use the correct method: localS.setItem()
localS.setItem("math_mentor_chats", json.dumps(data_to_save))