import streamlit as st
from streamlit_chat import message
import requests
from gpt_index import SimpleDirectoryReader,GPTListIndex,GPTSimpleVectorIndex,LLMPredictor,PromptHelper,ServiceContext
from langchain import OpenAI
import sys
import os

st.set_page_config(
    page_title="Chatbot",
    page_icon=":robot_face:"
)

def api_key_callback():
    # api_key = st.session_state.api_key_input
    os.environ["OPENAI_API_KEY"] = st.session_state.api_key_input
    # print(st.session_state.api_key_input)

def api_form():
    with st.form(key="search_form"):
        st.subheader("API Key")
        st.text_input("Enter your OpenAI API key", key="api_key_input")
        st.form_submit_button(label="Submit", on_click=api_key_callback)
        st.caption('''
            This api key wil not be stored and deleted at the end of your web session.
            For more details, you can check out the 
            [Source Code](https://huggingface.co/spaces/darthPanda/romeo_and_juliet_chatbot_with_gptIndex/tree/main). 
            If you want to be extra careful, you can delete your 
            [OpenAI API key](https://platform.openai.com/account/api-keys) 
            after using this space.
            ''')

with st.sidebar:
    st.title("Romeo and Juliet Chatbot")
    st.markdown('''
                <div>
                This gpt3 based chatbot has been 
                trained on famous play <b>Romeo and Juliet</b> by William Shakespeare
                </div>
                ''', unsafe_allow_html=True)
    api_form()
    # st.markdown("[Book Link](https://www.gutenberg.org/ebooks/1513)")
    st.markdown(
    "<div style='position: fixed; bottom: 0;'>Created by Taaha Bajwa</div>",
    unsafe_allow_html=True,
    )

st.markdown('Ask me any question from Romeo and Juliet')

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

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

def query(payload):
    vIndex = GPTSimpleVectorIndex.load_from_disk('vectorIndex.json')
    response = vIndex.query(payload,response_mode='compact')
    # print(str(response))
    return (str(response)).lstrip()

def get_text():
    input_text = st.text_input("You: ","", key="input")
    return input_text 


user_input = get_text()

if user_input:
    try:
        output = query(user_input)
        st.session_state.past.append(user_input)
        st.session_state.generated.append(output)
    except:
        st.error('Invalid API key', icon="🚨")

if st.session_state['generated']:
    for i in range(len(st.session_state['generated'])-1, -1, -1):
        message(st.session_state["generated"][i], key=str(i))
        message(st.session_state['past'][i], is_user=True, key=str(i) + '_user')