import streamlit as st

from langchain.chains import ConversationalRetrievalChain
from langchain.chat_models import ChatOpenAI
from langchain.memory import ConversationBufferMemory
from langchain.memory.chat_message_histories import StreamlitChatMessageHistory

from document_retriever import configure_retriever
from calback_handler import PrintRetrievalHandler, StreamHandler
from chat_profile import ChatProfileRoleEnum

# configs
LLM_MODEL_NAME = "gpt-3.5-turbo"

st.set_page_config(
    page_title=":books: InkChatGPT: Chat with Documents",
    page_icon="📚",
    initial_sidebar_state="collapsed",
    menu_items={
        "Get Help": "https://x.com/vinhnx",
        "Report a bug": "https://github.com/vinhnx/InkChatGPT/issues",
        "About": "InkChatGPT is a Streamlit application that allows users to upload PDF documents and engage in a conversational Q&A with a language model (LLM) based on the content of those documents.",
    },
)

# Setup memory for contextual conversation
msgs = StreamlitChatMessageHistory()

with st.container():
    col1, col2 = st.columns([0.2, 0.8])
    with col1:
        st.image(
            "./assets/large_icon.png", use_column_width="always", output_format="PNG"
        )
    with col2:
        st.header(":books: InkChatGPT")
        st.write("**Chat** with Documents")
        st.caption("Supports PDF, TXT, DOCX, EPUB • Limit 200MB per file")

chat_tab, documents_tab, settings_tab = st.tabs(["Chat", "Documents", "Settings"])
with settings_tab:
    openai_api_key = st.text_input("OpenAI API Key", type="password")
    if len(msgs.messages) == 0 or st.button("Clear message history"):
        msgs.clear()
        msgs.add_ai_message("How can I help you?")

with documents_tab:
    uploaded_files = st.file_uploader(
        label="Select files",
        type=["pdf", "txt", "docx", "epub"],
        accept_multiple_files=True,
    )

with chat_tab:
    if uploaded_files:
        result_retriever = configure_retriever(uploaded_files)

        memory = ConversationBufferMemory(
            memory_key="chat_history", chat_memory=msgs, return_messages=True
        )

        # Setup LLM and QA chain
        llm = ChatOpenAI(
            model_name=LLM_MODEL_NAME,
            openai_api_key=openai_api_key,
            temperature=0,
            streaming=True,
        )

        chain = ConversationalRetrievalChain.from_llm(
            llm, retriever=result_retriever, memory=memory, verbose=False
        )

        avatars = {
            ChatProfileRoleEnum.Human: "user",
            ChatProfileRoleEnum.AI: "assistant",
        }

        for msg in msgs.messages:
            st.chat_message(avatars[msg.type]).write(msg.content)

if not openai_api_key:
    st.caption("🔑 Add your **OpenAI API key** on the `Settings` to continue.")

if user_query := st.chat_input(
    placeholder="Ask me anything!", disabled=(not openai_api_key)
):
    st.chat_message("user").write(user_query)

    with st.chat_message("assistant"):
        retrieval_handler = PrintRetrievalHandler(st.empty())
        stream_handler = StreamHandler(st.empty())
        response = chain.run(user_query, callbacks=[retrieval_handler, stream_handler])