import streamlit as st
import os
from langchain_groq import ChatGroq
from langchain_community.agent_toolkits.load_tools import load_tools
from langchain_experimental.tools import PythonREPLTool
from langgraph.prebuilt import create_react_agent

# 1. Page Configuration
st.set_page_config(page_title="AI Math Agent", page_icon="🧮")
# Centered Title
st.markdown("<h1 style='text-align: center;'>MATH CHAT BOT 🧮</h1>", unsafe_allow_html=True)

# Centered Description
st.markdown("<p style='text-align: center;'>I'm here to help you solve your math homework step-by-step!</p>", unsafe_allow_html=True)

# 2. Secret Retrieval
api_key = os.getenv("GROQ_API_KEY")

if not api_key:
    st.warning("⚠️ Please add your GROQ_API_KEY to the Space Secrets in Settings.")
    st.stop()

# 3. Initialize the Brain (LLM)
# We use temperature=0 for mathematical precision
llm = ChatGroq(
    model_name="llama-3.3-70b-versatile", 
    temperature=0,
    groq_api_key=api_key
)

# ... (Keep your imports and LLM initialization the same)

# 4. Load Math Tools
# Load the tools FIRST
# Create the Python Tool
python_tool = PythonREPLTool()

# Load the standard math tools
math_tools = load_tools(["llm-math"], llm=llm)

# Combine them into a single list for the agent
tools = [python_tool] + math_tools
from langchain_core.messages import SystemMessage

system_message = SystemMessage(
    content=(
        "You are a specialized Mathematics Assistant. Your ONLY task is to solve problems "
        "within these branches: Arithmetic, Algebra, Geometry, Calculus, Analysis, "
        "Statistics & Probability, Trigonometry, Number Theory, and Topology.\n\n"
        "GUIDELINES:\n"
        "1. Use the 'python_repl' tool for complex formulas, symbolic algebra (SymPy), "
        "or statistical analysis (SciPy/NumPy).\n"
        "2. Use 'llm-math' for simple arithmetic.\n"
        "3. If a user asks about history, coding (non-math), or any non-mathematical topic, "
        "politely refuse.\n"
        "4. Always show your work or explain the mathematical steps taken."
    )
)

# Re-create the agent with the new tools and instructions
agent_executor = create_react_agent(llm, tools)

# 5. Chat Interface Setup
if "messages" not in st.session_state:
    st.session_state.messages = []

# Display conversation history
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

# 6. User Input Logic
if prompt := st.chat_input("Ask me a complex math question!"):
    # Add user message to history
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    # Generate Agent Response
    with st.chat_message("assistant"):
        with st.spinner("Thinking..."):
            try:
                # ALL lines below 'try' must be indented 4 spaces further
                # We pass the system_message rules directly here
                response = agent_executor.invoke({
                    "messages": [system_message, ("user", prompt)]
                })
                
                final_answer = response["messages"][-1].content
                
                st.markdown(final_answer)
                st.session_state.messages.append({"role": "assistant", "content": final_answer})
            
            except Exception as e:
                st.error(f"An error occurred: {e}")