import streamlit as st

from app import agent  # your CodeAgent instance
from Gradio_UI import stream_to_gradio  # generator yielding gr.ChatMessage objects

st.set_page_config(page_title="CodeAgent Chat", layout="wide")
st.title("CodeAgent Chat (Streamlit)")

# Initialize session state for chat history.
if "chat_history" not in st.session_state:
    st.session_state.chat_history = []

# Reset Chat button.
if st.button("Reset Chat"):
    st.session_state.chat_history = []
    st.rerun()  # Rerun the script to update the UI immediately.

# Display chat history using Streamlit's chat message container.
for chat in st.session_state.chat_history:
    st.chat_message(chat["role"]).write(chat["content"])

# Use st.chat_input for user message entry.
user_input = st.chat_input("Type your message here")
if user_input:
    # Append and display the user's message.
    st.session_state.chat_history.append({"role": "user", "content": user_input})
    st.chat_message("user").write(user_input)

    # Stream the agent responses.
    # The generator yields gr.ChatMessage objects that have 'role' and 'content' attributes.
    for msg in stream_to_gradio(agent, user_input, reset_agent_memory=False):
        role = msg.role
        # Convert non-string content to a string as needed.
        content = msg.content if isinstance(msg.content, str) else str(msg.content)
        st.session_state.chat_history.append({"role": role, "content": content})
        st.chat_message(role).write(content)