import streamlit as st
from agents import investigator, router
# from agents_openai_fc import investigator

st.title('Sonic Cyber Assistant')
st.caption("🚀 A cyber assistant chatbot for threat hunting")

if "messages" not in st.session_state:
    st.session_state["messages"] = [{"role": "assistant", "content": "How can I help you?"}]

for msg in st.session_state.messages:
    st.chat_message(msg["role"]).write(msg["content"])

if prompt := st.chat_input():

    st.session_state.messages.append({"role": "user", "content": prompt})
    st.chat_message("user").write(prompt)
    response = router.invoke(prompt)
    msg = response["output"]
    print("Title : ", router.generate_title(prompt))

    # print(response['intermediate_steps'])
    
    # If there's an existing assistant message, update it with the new response
    if st.session_state.messages[-1]["role"] == "assistant":
        st.session_state.messages[-1]["content"] = msg
    else:
        st.session_state.messages.append({"role": "assistant", "content": msg})

    st.chat_message("assistant").write(msg)

# Add a button to regenerate response
if st.button("Regenerate Response"):
    st.session_state.messages.pop()  # Remove the latest assistant message
    prompt = st.session_state.messages[-1]["content"]  # Retrieve user's last input
    with st.spinner("Searching..."):
        response = router.invoke(prompt)
        msg = response["output"]
        
        # If there's an existing assistant message, update it with the new response
        if st.session_state.messages[-1]["role"] == "assistant":
            st.session_state.messages[-1]["content"] = msg
        else:
            st.session_state.messages.append({"role": "assistant", "content": msg})

        st.chat_message("assistant").write(msg)

if st.button("Clear Chat"):
    st.session_state.messages = [{"role": "assistant", "content": "How can I help you?"}]
    router.memory.clear()
    for msg in st.session_state.messages:
        st.chat_message(msg["role"]).write(msg["content"])
    st.success("Chat cleared!")