import os

import openai
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
        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