import streamlit as st
from main import summarize_text, summarizer

st.set_page_config(page_title="TextSum AI", page_icon="📝", layout="centered")

st.title("📝 TextSum AI")
st.caption("Paste your text below and get a concise summary.")

# Clear memory button
if st.button("🗑️ Clear Memory"):
    summarizer.clear_memory()
    st.session_state.chat_history = []
    st.success("Memory cleared!")

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

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

# Input
user_input = st.chat_input("Paste your text or ask a follow-up question...")

if user_input:
    # Display user message
    with st.chat_message("user"):
        st.markdown(user_input)

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

    # Get summary
    with st.spinner("Summarizing..."):
        result = summarize_text(user_input)

    # Display AI response
    with st.chat_message("assistant"):
        st.markdown(result)

    st.session_state.chat_history.append({"role": "assistant", "content": result})