import streamlit as st
from model import Web_qa
from color import print_colorful_msg
import time

def chatbot_response(user_input):  # Open AI GPT-3 and Google T5
    start_time = time.time()
    print_colorful_msg(f"User Input: {user_input}", color='green')
    res = Web_qa(user_input)
    response = res['result']
    metadata = [i.metadata for i in res.get("source_documents", [])]
    end_time = time.time()
    response_time = end_time - start_time
    print_colorful_msg(f"Response Time: {response_time} seconds", color='yellow')
    return response, metadata, res.get('source_documents', [])

# def chatbot_response(user_input):  # Hugging Face  Misral model
#     start_time = time.time()
#     print_colorful_msg(f"User Input: {user_input}", color='green')
#     res = Web_qa(user_input)
#     print_colorful_msg(f"Response: {res}", color='blue')
#     response = res["result"].split(": Let me think about it...")[-1]
#     metadata = [i.metadata for i in res.get("source_documents", [])]
#     end_time = time.time()
#     response_time = end_time - start_time
#     print_colorful_msg(f"Response Time: {response_time} seconds", color='yellow')
#     return response, metadata, res.get('source_documents', [])

def main():
    print_colorful_msg("Starting chatbot main function ...", color='blue')
    st.title("ATrad Chatbot")

    user_input = st.text_input("Enter your message:")
    if st.button("Send"):
        response, metadata, source_documents = chatbot_response(user_input)
        st.text_area("Chatbot Response:", value=response, height=200)
        
        st.markdown("### Source Documents:")
        if metadata:
            for value in metadata:
                st.write(f"{value}")
        else:
            st.write("No metadata available.")
        
        # st.markdown("### Source Documents:")
        if source_documents:
            for doc in source_documents:
                pass
                # st.write(doc)
        else:
            st.write("No source documents found.")
        
        print_colorful_msg("End Chatbot Response...", color='blue')

if __name__ == "__main__":
    main()