import streamlit as st
from QnA import Q_A
import re,time
from QnA import get_hugging_face_model  , summarize ,get_groq_model


def summarize_data(documents,api_key):
    if api_key.startswith('gsk'):
        llm  = get_groq_model(api_key)
    else:
        llm= get_hugging_face_model(api_key=api_key)
    summary = summarize(documents,llm)
    return summary
    

def QA_Bot(vectorstore,API_KEY,documents):
    summary_response = None
    st.title("Q&A Bot")
    # Initialize chat history
    if "messages" not in st.session_state:
        st.session_state.messages = []

    # 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"])
    summary_response = summarize_data(documents,API_KEY)
    print(summary_response)


    # React to user input
    if prompt := st.chat_input("What is up?"):
        # Display user message in chat message container
        st.chat_message("user").markdown(prompt)
        # Add user message to chat history
        st.session_state.messages.append({"role": "user", "content": prompt})
        
        ai_response = Q_A(vectorstore,prompt,API_KEY)
        response = f"Echo: {ai_response}"
        # Display assistant response in chat message container
        with st.chat_message("assistant"):
            message_placeholder = st.empty()
            full_response = ""
            for chunk in re.split(r'(\s+)', response):
                full_response += chunk + " "
                time.sleep(0.01)

                # Add a blinking cursor to simulate typing
                message_placeholder.markdown(full_response + "▌")
        # Add assistant response to chat history
        st.session_state.messages.append({"role": "assistant", "content": full_response})