import os

import requests
import streamlit as st
from streamlit_chat import message

from models import bloom
from utils.util import *

st.title("Welcome to RegBotBeta")
st.header("Powered by `LlamaIndex🦙` and `OpenAI API`")

if 'messages' not in st.session_state:
    st.session_state.messages = []

api_key = st.text_input("Enter your OpenAI API key here:", type="password")
if api_key:
    resp = validate(api_key)
    if ("error" in resp.json()):
        st.info("Invalid Token! Try again.")
    else:
        st.info("Success")
        os.environ["OPENAI_API_KEY"] = api_key
        index = create_index(bloom)

st.write("---")
input_text = st.text_area("Ask your question")

if input_text is not None:
    if st.button("Ask"):
        st.session_state.messages.append(('User', input_text))
        with st.spinner("Processing your query..."):
            bot_response = get_response(index, input_text)
        print("bot: ", bot_response)
        st.session_state.messages.append(('Bot', bot_response))

# Display previous messages
msg_key = 0
for sender, msg in st.session_state.messages[::-1]:
    is_user = sender == "User"
    message(str(msg), is_user, key=str(msg_key)+f'_{sender}')
    msg_key += 1