import os

import requests
import streamlit as st

from models import bloom

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

def validate(token: str):
    api_endpoint = "https://api.openai.com/v1/chat/completions"
    api_key = token

    headers = {
        "Content-Type" : "application/json",
        "Authorization": f"Bearer {api_key}"
    }

    messages = [
        {"role": "user", "content": "Say this is a test!"}
    ]

    data = {
        "model": "gpt-3.5-turbo",
        "messages": messages
    }

    response = requests.post(api_endpoint, json=data, headers=headers)
    return response

def create_index():
    index = bloom.initialize_index("bloomLlama")
    return index

def get_response(vector_index, query_str):
    query_engine = vector_index.as_query_engine()
    response = query_engine.query(query_str)
    return response


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("Your API Token is incorrect! Try again.")
    else:
        os.environ["OPENAI_API_KEY"] = api_key
        index = create_index()

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

if input_text is not None:
    if st.button("Ask"):
        st.info("Your query: \n" + input_text)
        with st.spinner("Processing your query..."):
            response = get_response(index, input_text)
            print(response)

        st.success(response)

        st.write("---")