# frontend.py
import streamlit as st
import asyncio
from backend import get_agent, MCPAgent
import time

st.set_page_config(
    page_title="MCP Agent Chat",
    page_icon="🤖",
    layout="wide",
    initial_sidebar_state="expanded"
)

st.markdown("""
<style>
    .stChatMessage {
        padding: 1rem;
        border-radius: 0.5rem;
        margin-bottom: 1rem;
    }
    .user-message {
        background-color: #e3f2fd;
    }
    .assistant-message {
        background-color: #f5f5f5;
    }
    .sidebar-info {
        padding: 1rem;
        background-color: #f0f2f6;
        border-radius: 0.5rem;
        margin-bottom: 1rem;
    }
</style>
""", unsafe_allow_html=True)

# Initialize session state
if "messages" not in st.session_state:
    st.session_state.messages = []
    st.session_state.history = []  # For agent history
    st.session_state.agent_initialized = False
    st.session_state.available_tools = []  # Add this to ensure iterable default
    st.session_state.pending_query = None  # For example queries

# Sidebar
with st.sidebar:
    st.title("🤖 MCP AI Algo Trade App Demo")
    st.markdown("### ℹ️ About")
    st.markdown("""
    This chat interface connects to:
    - **Math Server**: Basic arithmetic operations
    - **Stock Server**: Mock Real-time market data
    
    The agent uses LangChain and MCP to intelligently route your queries to the appropriate tools.
    """)
    st.markdown("---")
    
    if st.session_state.agent_initialized:
        st.success("✅ Agent Connected")
        
        st.markdown("### 🛠️ Available Tools")
        available_tools = st.session_state.get("available_tools", [])
        for tool in available_tools:
            st.markdown(f"• `{tool}`")
    else:
        st.info("⏳ Agent Initializing...")
    
    st.markdown("---")
    
    st.markdown("### 💡 Example Queries")
    example_queries = [
        "What's the price of AAPL?",
        "Show me the market summary",
        "Get news about TSLA",
        "Calculate 25 * 4",
        "What's 100 divided by 7?",
        "Add 456 and 789"
    ]
    
    for query in example_queries:
        if st.button(query, key=f"example_{query}"):
            st.session_state.pending_query = query
            st.rerun()
    
    st.markdown("---")
    
    if st.button("🗑️ Clear Chat", type="secondary"):
        st.session_state.messages = []
        st.session_state.history = []
        st.rerun()
    
    

# Main chat interface
st.title("💬 MCP Agent Chat")
st.markdown("Ask me about stocks, math calculations, or general questions!")

# Initialize agent asynchronously
async def initialize_agent():
    """Initialize the agent if not already done"""
    if not st.session_state.agent_initialized:
        agent = get_agent()
        with st.spinner("🔧 Initializing MCP servers..."):
            try:
                tools = await agent.initialize()
                st.session_state.available_tools = tools or []  # Guard against None
                st.session_state.agent_initialized = True
                st.rerun()  # Add this to refresh UI after init
                return True
            except Exception as e:
                st.error(f"Failed to initialize agent: {str(e)}")
                st.info("Please make sure the stock server is running: `python stock_server.py`")
                return False
    return True

async def process_user_message(user_input: str):
    """Process the user's message and get response from agent"""
    agent = get_agent()
    

    st.session_state.history.append({"role": "user", "content": user_input})
    
    try:

        response = await agent.process_message(user_input, st.session_state.history)
        

        st.session_state.history.append({"role": "assistant", "content": response})
        
        return response
    except Exception as e:
        return f"Error: {str(e)}"


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


if st.session_state.pending_query:
    query = st.session_state.pending_query
    st.session_state.pending_query = None  # Clear it
    
    st.session_state.messages.append({"role": "user", "content": query})
    
    async def process_example():
        if not st.session_state.agent_initialized:
            if not await initialize_agent():
                return "Failed to initialize agent. Please check the servers."
        return await process_user_message(query)
    
    with st.spinner("Processing..."):
        response = asyncio.run(process_example())
    
    st.session_state.messages.append({"role": "assistant", "content": response})
    
    st.rerun()

if prompt := st.chat_input("Type your message here..."):
    st.session_state.messages.append({"role": "user", "content": prompt})
    
    with st.chat_message("user"):
        st.markdown(prompt)
    
    with st.chat_message("assistant"):
        message_placeholder = st.empty()
        
        async def get_response():
            if not st.session_state.agent_initialized:
                if not await initialize_agent():
                    return "Failed to initialize agent. Please check the servers."
            
            return await process_user_message(prompt)
        
        with st.spinner("Thinking..."):
            response = asyncio.run(get_response())
        
        message_placeholder.markdown(response)
        st.session_state.messages.append({"role": "assistant", "content": response})

if not st.session_state.agent_initialized:
    with st.spinner("🔧 Initializing agent..."):
        asyncio.run(initialize_agent())

st.markdown("---")
st.markdown(
    """
    <div style='text-align: center; color: #666;'>
    Powered by LangGraph, LangChain, MCP, and Hugging Face
    Developed by Lorentz Yeung
    </div>
    """,
    unsafe_allow_html=True
)