# gradio imports
import gradio as gr
import os
import time

# Imports
import os

import openai
from langchain.chains import ConversationalRetrievalChain

from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.chat_models import ChatOpenAI
from langchain.text_splitter import CharacterTextSplitter
from langchain.vectorstores import Chroma
from langchain.document_loaders import TextLoader

from langchain.memory import ConversationBufferMemory
from langchain.chat_models import ChatOpenAI

css="""
#col-container {max-width: 700px; margin-left: auto; margin-right: auto;}
"""

title = """
<div style="text-align: center;max-width: 700px;">
    <h1>Chat about Dialogues • Games • AI • AI Regulation</h1>
    <p style="text-align: left;">Chat is built from:<br />
    This is a Dialogue (https://www.jonnyjohnson.com/this-is-a-dialogue)<br />
    Game-Making articles (https://dialogues-ai.github.io/papers/docs/ai_regulation/gamemaking)<br />
    As well as 25 blog posts contributed to BMC <br />
</div>
"""


prompt_hints = """
<div style="text-align: center;max-width: 700px;">
    <p style="text-align: left;">Some things you can ask:<br />
    Should I be worried about AIs?<br />
    How do we improve the games between AIs and humans?<br />
    What is a dialogue? <br />
    Do you agree that everything is language? <br />
</div>
"""

# from index import PERSIST_DIRECTORY, CalendarIndex
PERSIST_DIRECTORY = "chromadb"
# Create embeddings

# # create memory object
from langchain.memory import ConversationBufferMemory
memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)

def loading_pdf():
    return "Loading..."

def loading_database(open_ai_key):
    if open_ai_key is not None:
        if os.path.exists(PERSIST_DIRECTORY):
            embeddings = OpenAIEmbeddings(openai_api_key=open_ai_key)
            docs_retriever = Chroma(persist_directory=PERSIST_DIRECTORY, embedding_function=embeddings)

            global qa_chain
            qa_chain = ConversationalRetrievalChain.from_llm(ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0.0, openai_api_key=open_ai_key),  
                            retriever=docs_retriever.as_retriever(), 
                            memory=memory, 
                            return_source_documents=False
                            )
            return "Ready"
    else:
        return "You forgot OpenAI API key"

def add_text(history, text):
    history = history + [(text, None)]
    return history, ""


def bot(history):
    response = infer(history[-1][0], history)
    history[-1][1] = ""
    for character in response:     
        history[-1][1] += character
        time.sleep(0.05)
        yield history
    

def infer(question, history):
    res = []
    for human, ai in history[:-1]:
        pair = (human, ai)
        res.append(pair)
    
    chat_history = res
    query = question
    result = qa_chain({"question": query, "chat_history": chat_history})
    return result["answer"]

def update_message(question_component, chat_prompts):
    question_component.value = chat_prompts.get_name()
    return None

with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.HTML(title)
        with gr.Column():
            with gr.Row():
                openai_key = gr.Textbox(label="OpenAI API key", type="password")
                submit_api_key = gr.Button("Submit")
            with gr.Row():
                langchain_status = gr.Textbox(label="Status", placeholder="", interactive=False)
        
        chatbot = gr.Chatbot([], elem_id="chatbot").style(height=350)
        question = gr.Textbox(label="Question", placeholder="Type your question and hit Enter ")
        submit_btn = gr.Button("Send Message")
        gr.HTML(prompt_hints)
        
    submit_api_key.click(loading_database, inputs=[openai_key], outputs=[langchain_status], queue=False)
    # demo.load(loading_database, None, langchain_status)
    question.submit(add_text, [chatbot, question], [chatbot, question]).then(
        bot, chatbot, chatbot
    )
    submit_btn.click(add_text, [chatbot, question], [chatbot, question]).then(
        bot, chatbot, chatbot)

demo.queue(concurrency_count=2, max_size=20).launch()