import altair as alt
import streamlit as st
import pandas as pd
import os
import zipfile
from PIL import Image
import sympy as sp
import time
import tempfile

# Page config
st.set_page_config(page_title="🎓 Smart Math Teacher", layout="centered")

# Custom CSS
st.markdown("""
    <style>
    .fun-title { font-size: 40px; color: #ff3399; text-align: center; font-family: 'Comic Sans MS', cursive; }
    .question-box { border: 4px dotted #ffcc00; padding: 20px; border-radius: 20px; background-color: #fff7e6; font-size: 20px; }
    .stButton > button { font-size: 18px; background-color: #00cc99; color: white; border-radius: 10px; padding: 10px; }
    .stTextInput > div > input { font-size: 18px; }
    
    /* Guide styling */
    .guide-box {
        background-color: #e8f5e8;
        padding: 15px;
        border-radius: 10px;
        border: 2px solid #4CAF50;
        margin: 10px 0;
    }
    </style>
""", unsafe_allow_html=True)

# Welcome title
st.markdown("<div class='fun-title'>🧠✨ Welcome to the Smart Math Teacher! ✨🧠</div>", unsafe_allow_html=True)

# USER GUIDE - Always visible at the top
st.markdown("""
<div class="guide-box">
<h3>📖 How to Use This App:</h3>

**1. Select Your Age Group** - Choose your learning level (4-6, 7-9, or 13-15 years)<br>
**2. Choose Math Category** - Pick what type of math to practice<br>
**3. Solve Questions** - Read each question and type your answer<br>
**4. Get Help** - Use hints or skip if needed<br>
**5. Learn** - See correct answers and step-by-step solutions
</div>
""", unsafe_allow_html=True)

# Initialize session state for temp directory
if 'temp_dir' not in st.session_state:
    st.session_state.temp_dir = tempfile.mkdtemp()

# Age group setup 
age_groups = {
    "4-6 Age Group": {"dataset": "src/Datase_of_4-6_Age_Group.xlsx", "zip_file": "src/Image_for_group_4-6.zip", "image_folder": "Image_for_group_4-6"},
    "7-9 Age Group": {"dataset": "src/Datase_of_7-9_Age_Group.xlsx", "zip_file": "src/Image_for_group_7-9.zip", "image_folder": "Image_for_group_7-9"},
    "13-15 Age Group": {"dataset": "src/Datase_of_13-15_Age_Group.xlsx", "zip_file": "src/Image_for_group_13-15.zip", "image_folder": "Image_for_group_13-15"},
}

selected_age_group = st.selectbox("🧒 Select your Age Group:", list(age_groups.keys()))

# Initialize session
if "session_initialized" not in st.session_state or st.session_state.age_group != selected_age_group:
    st.session_state.age_group = selected_age_group
    st.session_state.category = None
    st.session_state.question_index = 0
    st.session_state.show_answer = False
    st.session_state.show_steps = False
    st.session_state.session_initialized = True

# Load dataset
group_info = age_groups[selected_age_group]
dataset_path = group_info["dataset"]
zip_path = group_info["zip_file"]
image_folder = group_info["image_folder"]

# Create image folder in temp directory
temp_image_folder = os.path.join(st.session_state.temp_dir, image_folder)
os.makedirs(temp_image_folder, exist_ok=True)

# Extract images if zip exists
try:
    if os.path.exists(zip_path):
        with zipfile.ZipFile(zip_path, "r") as zip_ref:
            zip_ref.extractall(temp_image_folder)
except Exception as e:
    st.warning(f"Error extracting images: {e}")

if not os.path.exists(dataset_path):
    st.error(f"Dataset not found: {dataset_path}")
    st.stop()

try:
    df = pd.read_excel(dataset_path)
    df['category'] = df['category'].astype(str).str.strip()
except Exception as e:
    st.error(f"Error loading dataset: {e}")
    st.stop()

# Category selection
categories = sorted(df['category'].dropna().unique())
selected_category = st.selectbox("📚 Choose a Math Category:", options=categories)

# Update session if category changes
if st.session_state.category != selected_category:
    st.session_state.category = selected_category
    st.session_state.question_index = 0
    st.session_state.show_answer = False
    st.session_state.show_steps = False
    st.rerun()

# Filter questions by selected category only
subset_df = df[df['category'] == selected_category].reset_index(drop=True)

if not subset_df.empty and st.session_state.question_index < len(subset_df):
    question = subset_df.iloc[st.session_state.question_index]
    progress = int((st.session_state.question_index / len(subset_df)) * 100)
    st.progress(progress)

    st.markdown(f"<div class='question-box'>📘 <b>Question {st.session_state.question_index + 1}:</b><br><br>{question['problem']}</div>", unsafe_allow_html=True)

    # Display image with error handling
    if pd.notna(question.get('image')):
        image_name = str(question['image']).strip()
        image_found = False
        for root, _, files in os.walk(temp_image_folder):
            for file in files:
                if file.lower().startswith(image_name.lower()) or os.path.splitext(file)[0].lower() == image_name.lower():
                    try:
                        st.image(Image.open(os.path.join(root, file)), use_column_width=True)
                        image_found = True
                        break
                    except:
                        pass
        if not image_found:
            st.warning("❌ Image not found.")

    user_ans = st.text_input("📝 Your Answer:", key=f"ans_{st.session_state.question_index}")

    if st.button("✅ Submit Answer"):
        if str(user_ans).strip().lower() == str(question['answer']).strip().lower():
            st.success("🎉 Correct! Well done!")
            st.balloons()
            time.sleep(2)
            st.session_state.question_index += 1
            st.session_state.show_answer = False
            st.session_state.show_steps = False
            st.rerun()
        else:
            st.error("❌ Try again or view the correct answer below.")
            st.session_state.show_answer = True
            st.session_state.show_steps = False

    if st.session_state.show_answer:
        st.info(f"✅ Correct Answer: **{question['answer']}**")
        if selected_age_group in ["7-9 Age Group", "13-15 Age Group"]:
            if st.button("🔍 Show Steps"):
                st.session_state.show_steps = True
        if st.session_state.show_steps and pd.notna(question.get("steps", None)):
            st.success(f"### 🪄 Steps:\n{question['steps']}")

    if st.button("⏭️ Skip"):
        st.session_state.question_index += 1
        st.session_state.show_answer = False
        st.session_state.show_steps = False
        st.rerun()
elif subset_df.empty:
    st.warning("⚠️ No questions available in this category. Try another one.")
else:
    st.success("🏁 You've completed all questions in this category!")

# Cleanup function
def cleanup():
    import shutil
    if 'temp_dir' in st.session_state and os.path.exists(st.session_state.temp_dir):
        shutil.rmtree(st.session_state.temp_dir)

import atexit
atexit.register(cleanup)