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

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

def main_RAG_ui():
    use_chunk = st.sidebar.slider(
        'Chunk Size',
        500, 2000, (1000)
    )
    use_temperature = st.sidebar.slider(
        'Temperature 🦄',
        0.0, 1.0, (0.1))
    use_compression = st.checkbox("Compression🛠️(on uploaded document)", value=False)
    use_ddg_search = st.checkbox("Search on DuckDuckGO🦆(does not use document)", value=False)

    
    CONV_CHAIN = config_retrieval_chain(
        uploaded_files,
        use_compression=use_compression,
        use_chunksize=use_chunk,
        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:
                params = {
                    "question": user_query,
                    "chat_history": MEMORY.chat_memory.messages,
                }
                response = CONV_CHAIN.run(params, callbacks=[stream_handler])
                if response:
                    container.markdown(response)


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

uploaded_files = st.sidebar.file_uploader(
    label="Upload a file🐣",
    type=list(DocumentLoader.supported_extensions.keys()),
    accept_multiple_files=True
)

if not uploaded_files:
    st.info("Upload a file to start🐣")
    st.stop()

main_RAG_ui()