import os
import time
import gradio as gr
import openai

from langdetect import detect 
from gtts import gTTS
from pdfminer.high_level import extract_text

#any vector server should work, trying pinecone first
import pinecone

#langchain part
import spacy
import tiktoken
from langchain.llms import OpenAI
from langchain.text_splitter import SpacyTextSplitter
from langchain.document_loaders import TextLoader
from langchain.document_loaders import DirectoryLoader
from langchain.indexes import VectorstoreIndexCreator
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import Pinecone
import markdown


openai.api_key = os.environ['OPENAI_API_KEY']
pinecone_key = os.environ['PINECONE_API_KEY_AMD']
pinecone_environment='us-west1-gcp-free'


user_db = {os.environ['username1']: os.environ['password1'], os.environ['username2']: os.environ['password2'], os.environ['username3']: os.environ['password3']}

messages = [{"role": "system", "content": 'You are a helpful assistant.'}]
errors = []
error_recorded = 1



#load up spacy

nlp = spacy.load("en_core_web_sm")



def init_pinecone():
    pinecone.init(api_key=pinecone_key, environment=pinecone_environment)
    return




def process_file(index_name, dir):

    init_pinecone()

    #using openai embedding hence dim = 1536
    pinecone.create_index(index_name, dimension=1536, metric="cosine")
    #time.sleep(5)
    
    embeddings = OpenAIEmbeddings(openai_api_key=os.environ['OPENAI_API_KEY'])
    splter = SpacyTextSplitter(chunk_size=1000,chunk_overlap=200)

    for doc in dir:
        loader = TextLoader(doc.name , encoding='utf8')
        content = loader.load()
        split_text = splter.split_documents(content)
        for text in split_text:
        	Pinecone.from_documents([text], embeddings, index_name=index_name)

    #pipeline='zh_core_web_sm'
    

    return 


def list_pinecone():
    init_pinecone()
    return pinecone.list_indexes()


def show_pinecone(index_name):
    init_pinecone()
    #return pinecone.describe_index(index_name)
    index = pinecone.Index(index_name)
    stats = index.describe_index_stats()
    return stats



def delete_pinecone(index_name):
    init_pinecone()
    pinecone.delete_index(index_name)
    return


# Record feed back

def not_in_error():
    global error_recorded
    if(error_recorded):
        return
    else:
        global messages
        error_recorded = 1
    error = ("not_in_error", messages)
    errors.append(error)
    return "Thank you, the question has been marked as not in ROCm."


def not_found_error():
    global error_recorded
    if(error_recorded):
        return
    else:
        global messages
        error_recorded = 1
    error = ("not_found_error", messages)
    errors.append(error)
    return "Thank you, the context has been marked as not found in Vector Server."



def llm_error():
    global error_recorded
    if(error_recorded):
        return
    else:
        global messages
        error_recorded = 1
    error = ("not_found_error", messages)
    errors.append(error)
    return "Thank you, the LLM error has been recorded correctly."

def list_erros():
    global errors
    result = '\n\n\n'.join([f'({x}, {y})' for x, y in errors])
    return result

def clear_erros():
    global errors
    errors = []
    return "Warning, you just delete all the store errors!"









def roleChoice(role):
    global messages
    messages = [{"role": "system", "content": role}]
    return "role:" + role






def talk2file(index_name, text):
    #disable the global message
    global messages
    global error_recorded
    error_recorded = 0
    messages = [{"role": "system", "content": 'You are a helpful assistant.'}]
    
    #same as filesearch
    init_pinecone()
    embeddings = OpenAIEmbeddings(openai_api_key=os.environ['OPENAI_API_KEY'])
    docsearch = Pinecone.from_existing_index(index_name, embeddings)
    docs = docsearch.similarity_search(text)

    
    prompt = text + ", based on the following context: \n\n" 
    qwcontext = prompt + docs[0].page_content
    messages.append({"role": "user", "content": qwcontext})

    response = openai.ChatCompletion.create(model="gpt-3.5-turbo", messages=messages)

    system_message = response["choices"][0]["message"]
    #messages.append(system_message)

    #chats = ""
    #for msg in messages:
    #    if msg['role'] != 'system':
    #        chats += msg['role'] + ": " + msg['content'] + "\n\n"

    #Title1 = '<h2 style="background-color: yellow;"><b>User Question: </b></h2>'
    #User_Question = f'<div style="background-color: yellow; display: inline-block; word-wrap: break-word;">{prompt}</div>'
    Title2 = '<h2><b>Context Found: </b></h2>'
    context = docs[0].page_content
    #context = f'<span style="background-color: #ADD8E6; display: inline-block; word-wrap: break-word;">{context2html}</span>'
    #Title3 = '<h2 style="background-color: #90EE90;"><b>Ansewr: </b></h2>'
    answer = system_message["content"]


    
    return [context, answer]





def fileSearch(index_name, prompt):
    global messages

    init_pinecone()
    embeddings = OpenAIEmbeddings(openai_api_key=os.environ['OPENAI_API_KEY'])
    docsearch = Pinecone.from_existing_index(index_name, embeddings)
    docs = docsearch.similarity_search(prompt)

    return ["# Top1 context:\n" + docs[0].page_content, "# Top2 context:\n" + docs[1].page_content, "# Top3 context:\n" + docs[2].page_content]



def clear():
    global messages
    messages = [{"role": "system", "content": 'You are a helpful technology assistant.'}]
    return
    
def show():
    global messages
    chats = ""
    for msg in messages:
        if msg['role'] != 'system':
            chats += msg['role'] + ": " + msg['content'] + "\n\n"

    return chats


# feed back record
with gr.Blocks() as FeedBack:
    gr.Markdown("Record Feedback for the ROCm Usage Tutor, using the following three buttons to record three different errors:\n 1. ROCm-related contexts not included in the ROCm repo.\n 2. Context included in the repo, but the vector server fails to find it.\n 3. The LLM model fails to understand the context.")
    notin_btn = gr.Button("Not_In_Err")
    notin = gr.Textbox()
    notin_btn.click(fn=not_in_error, inputs=None, outputs=notin, queue=False)
    notfound_btn = gr.Button("Not_Found_Err")
    notfound = gr.Textbox()
    notfound_btn.click(fn=not_found_error, inputs=None, outputs=notfound, queue=False)
    llmerr_btn = gr.Button("LLM_Err")
    llmerr = gr.Textbox()
    llmerr_btn.click(fn=llm_error, inputs=None, outputs=llmerr, queue=False)
    listerr_btn = gr.Button("List_Errors")
    listerr = gr.Textbox()
    listerr_btn.click(fn=list_erros, inputs=None, outputs=listerr, queue=False)
    clearerr_btn = gr.Button("Clear_Errors")
    clearerr = gr.Textbox()
    clearerr_btn.click(fn=clear_erros, inputs=None, outputs=clearerr, queue=False)

    


with gr.Blocks() as chatHistory:
    gr.Markdown("Click the Clear button below to remove all the chat history.")
    clear_btn = gr.Button("Clear")
    clear_btn.click(fn=clear, inputs=None, outputs=None, queue=False)

    gr.Markdown("Click the Display button below to show all the chat history.")
    show_out = gr.Textbox()
    show_btn = gr.Button("Display")
    show_btn.click(fn=show, inputs=None, outputs=show_out, queue=False)


#pinecone tools
with gr.Blocks() as pinecone_tools: 
    pinecone_list = gr.Textbox()
    list = gr.Button(value="List") 
    list.click(fn=list_pinecone, inputs=None, outputs=pinecone_list, queue=False)

    pinecone_delete_name = gr.Textbox()
    delete = gr.Button(value="Delete") 
    delete.click(fn=delete_pinecone, inputs=pinecone_delete_name, outputs=None, queue=False)

    pinecone_show_name = gr.Textbox()
    pinecone_info = gr.Textbox()
    show = gr.Button(value="Show") 
    show.click(fn=show_pinecone, inputs=pinecone_show_name, outputs=pinecone_info, queue=False)



    


textbox = gr.inputs.Textbox(label="Vector Server Index Name: ", default="amd")
textbox2 = gr.inputs.Textbox(label="Vector Server Index Name: ", default="amd")
answerbox = gr.inputs.Textbox(label="Assistant answer")
contextbox = gr.Markdown(label="Context found")


contextbox1 = gr.Markdown(label="Top1 context")
contextbox2 = gr.Markdown(label="Top2 context")
contextbox3 = gr.Markdown(label="Top3 context")

role = gr.Interface(fn=roleChoice, inputs="text", outputs="text", description = "Choose your GPT roles, e.g. You are a helpful technology assistant.")
text = gr.Interface(fn=talk2file, inputs=[textbox, "text"], outputs=[contextbox, answerbox])

vector_server = gr.Interface(fn=process_file, inputs=["text", gr.inputs.File(file_count="directory")], outputs="text")

#audio = gr.Interface(fn=audioGPT, inputs=gr.Audio(source="microphone", type="filepath"), outputs="text")
#siri = gr.Interface(fn=siriGPT, inputs=gr.Audio(source="microphone", type="filepath"), outputs = "audio")
file = gr.Interface(fn=fileSearch, inputs=[textbox2, "text"], outputs=[contextbox1, contextbox2, contextbox3], description = "This tab shows the top three most related contexts in the repository.")
#demo = gr.TabbedInterface([role, text, file, vector_server, pinecone_tools, chatHistory], [ "roleChoice", "Talk2File", "FileSearch", "VectorServer", "PineconeTools", "ChatHistory"])

demo = gr.TabbedInterface([text, file, FeedBack], [ "ROCm Usage Tutor", "Top 3 Context", "Feedback"])

if __name__ == "__main__":
    demo.launch(enable_queue=False, auth=lambda u, p: user_db.get(u) == p,
        auth_message="This is not designed to be used publicly as it links to a personal openAI API. However, you can copy my code and create your own multi-functional ChatGPT with your unique ID and password by utilizing the 'Repository secrets' feature in huggingface.")
    #demo.launch()