import streamlit as st
import logging
from utils import MEMORY, DocumentLoader, check_password
from chat import config_noretrieval_chain
from streamlit.external.langchain import StreamlitCallbackHandler

logging.basicConfig(encoding="utf-8", level=logging.INFO)
LOGGER = logging.getLogger()

def main_chat_ui():
    use_temperature = st.sidebar.slider(
        'Temperature 🦄',
        0.0, 1.0, (0.1))
    use_ddg_search = st.checkbox("Search on DuckDuckGO🦆", value=False)
    
    CONV_CHAIN = config_noretrieval_chain(
        use_temperature=use_temperature,
        use_zeroshoot=use_ddg_search
    )
    if st.sidebar.button("Clear History🦭"):
        MEMORY.chat_memory.clear()
    if len(MEMORY.chat_memory.messages) == 0:
        st.chat_message("assistant").markdown("Ask me something🤖")
    
    avatars = {"human": "user", "ai": "assistant"}
    if user_query := st.chat_input(placeholder="Say something🐻"):
        st.chat_message("user").write(user_query)
        container = st.empty()
        stream_handler = StreamlitCallbackHandler(container)
        with st.chat_message("assistant"):
            if use_ddg_search:
                response = CONV_CHAIN.invoke(
                    {"input": user_query}, {"callbacks": [stream_handler]}
                )
                st.write(response["output"])
            else:
                response = CONV_CHAIN.run(user_query)
                if response:
                    container.markdown(response)


#if not check_password():
#    st.stop()
st.title("👻START CHAT👻")
main_chat_ui()