import streamlit as st
from src.buildgraph import run_workflow
import time

st.set_page_config(page_title="LawGPT")
col1, col2, col3 = st.columns([1,4,1])
with col2:
    st.image("assets/Black Bold Initial AI Business Logo.jpg")

st.markdown(
    """
    <style>
    .stApp, .ea3mdgi6{
      background-color:#000000;
    }
    div.stButton > button:first-child {
        background-color: #ffd0d0;
    }
    div.stButton > button:active {
        # background-color: #ff6262;
    }
    div[data-testid="stStatusWidget"] div button {
        display: none;
    }
    .reportview-container {
        margin-top: -2em;
    }
    #MainMenu {visibility: hidden;}
    .stDeployButton {display:none;}
    footer {visibility: hidden;}
    #stDecoration {display:none;}
    button[title="View fullscreen"]{
        visibility: hidden;
    }
    button:first-child{
        background-color : transparent !important;
    }
    </style>
    """,
    unsafe_allow_html=True,
)

st.title("AI Chatbot")

# Initialize chat history and thread_id
if "messages" not in st.session_state:
    st.session_state.messages = []
if "thread_id" not in st.session_state:
    st.session_state.thread_id = "streamlit_thread"

config = {"recursion_limit": 8, "configurable": {"thread_id": st.session_state.thread_id}}

# Display chat messages from history on app rerun
for message in st.session_state.messages:
    with st.chat_message(message['role']):
        st.markdown(message['content'])

# React to user input
if prompt := st.chat_input("What is your question?"):
    # Display user message in chat message container
    st.chat_message("user").markdown(prompt)
    user_message = {"role": "user", "content": prompt}
    # Add user message to chat history
    st.session_state.messages.append(user_message)

    response = run_workflow(prompt, config)
    response_content = response.get("generation", "I'm sorry, I couldn't generate a response.")

    
    # Display assistant response in chat message container
    with st.chat_message("assistant"):
        message_placeholder = st.empty()
        full_response = "⚠️ **_Note: Information provided may be inaccurate._** \n\n\n"
        for char in response_content:
            full_response += char
            
            time.sleep(0.03)
            message_placeholder.markdown(full_response + "|")
        message_placeholder.markdown(full_response)
        print(full_response)
    
    # Add assistant response to chat history
    st.session_state.messages.append({"role": "assistant", "content": full_response})

def reset_conversation():
    st.session_state.messages = []

st.button('Reset All Chat 🗑️', on_click=reset_conversation)