from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import Chroma
from langchain.chat_models import ChatOpenAI
from langchain.prompts.chat import (
    ChatPromptTemplate,
    SystemMessagePromptTemplate,
    HumanMessagePromptTemplate,
)
from langchain.chains import ChatVectorDBChain
import gradio as gr

# initialize variables
query = ""
chat_history = []
system_template = "You are a helpful assistant. Use the following pieces of context to \
answer the user's question. {context}"
temperature = 0
qa = None

# function called when user submits a question
def chatbot(input, history=[]):
    result = qa({"question": input, "chat_history": []})
    history.append((input, result["answer"]))
    return history, history


# function to set OpenAI key and initalize LangChain chain
def set_open_ai_key(input):
    print("Setting OpenAI key to ", input)
    # os.environ["OPENAI_API_KEY"] = input

    # read in embedddings from vector db store
    embeddings = OpenAIEmbeddings(openai_api_key=input)
    vectordb = Chroma(persist_directory="vector_store", embedding_function=embeddings)

    # initialize chatbot message templates
    messages = [
        SystemMessagePromptTemplate.from_template(system_template),
        HumanMessagePromptTemplate.from_template("{question}\Standalone question"),
    ]
    prompt = ChatPromptTemplate.from_messages(messages)

    # initialize chatbot chain
    global qa
    qa = ChatVectorDBChain.from_llm(
        ChatOpenAI(temperature=temperature, openai_api_key=input),
        vectordb,
        qa_prompt=prompt,
    )


# example questions user can choose from in Gradio UI to test the chatbot
examples = [
    "What does Biggie Smalls say about learning from other people's mistakes?",
    "What can I learn from Megan Quinn about reading?",
    "What does Scott Belsky say about optimizing what works?",
    "What is product market fit?",
    "What is the best way to lower CAC?",
    "What is the most important competitive advantage a business can have today?",
]

# initialize Gradio UI using chatbot UI
gr_interface = gr.Interface(
    fn=chatbot,
    inputs=["text", "state"],
    outputs=["chatbot", "state"],
    title="TrenBot",
    description="Q&A Chatbot for Tren Griffin's views on the \
             market, tech, and everything else",
    examples=examples,
    thumbnail="https://i0.wp.com/25iq.com/wp-content/uploads/2012/10/tren_griffin_crop.jpg",
)

# add OpenAI key textbox to Gradio UI
with gr_interface:
    openai_key_textbox = gr.Textbox(
        lines=1,
        placeholder="Paste your OpenAI key here and hit Enter.",
        type="password",
        label="Your OpenAI Key",
    )
    openai_key_textbox.submit(set_open_ai_key, inputs=openai_key_textbox)

# launch Gradio UI
gr_interface.launch()