import streamlit as st

# Import your agent and streaming function.
# For example, if your agent code is in agent_app.py:
from app import agent  # your CodeAgent instance
from Gradio_UI import (
    stream_to_gradio,
)  # re-use the generator that yields gr.ChatMessage

# (Optionally, if you want to avoid Gradio-specific types you can write your own streaming generator.)

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 = []


def display_chat():
    """Display the chat history in the app."""
    for message in st.session_state.chat_history:
        role = message.get("role", "assistant")
        content = message.get("content", "")
        if role == "user":
            st.markdown(f"**User:** {content}")
        else:
            st.markdown(f"**Assistant:** {content}")


# Main chat container.
chat_container = st.container()
with chat_container:
    display_chat()

# TODO: use `st.chat_input`

# User input.
user_input = st.text_input("Enter your message:", key="input_text")
if st.button("Send") and user_input:
    # Append the user message to the history.
    st.session_state.chat_history.append({"role": "user", "content": user_input})
    with chat_container:
        display_chat()

    # Stream the agent responses.
    # (Here we are reusing your existing streaming generator.
    #  Note that gr.ChatMessage objects have attributes "role" and "content".)
    placeholder = st.empty()  # if you want to update a placeholder
    for msg in stream_to_gradio(agent, user_input, reset_agent_memory=False):
        # Extract role and content.
        role = msg.role
        # For content that is not a plain string (e.g. images or audio), you might need to
        # add extra handling. Here we simply convert it to a string.
        content = msg.content if isinstance(msg.content, str) else str(msg.content)
        st.session_state.chat_history.append({"role": role, "content": content})
        with chat_container:
            display_chat()