import openai
from llama_index import SimpleDirectoryReader, GPTVectorStoreIndex, Prompt
from cnocr import CnOcr
import gradio as gr
import time

ocr = CnOcr()
history_max_len = 500
all_max_len = 2000

def get_embeddings_from_texts(openai_apikey,text):
    openai.api_key = openai_apikey
    response = openai.Embedding.create(
        input = text,
        model = "text-embedding-ada-002"
        
    )
    return reponse['data'][0]['embedding']

def doc_index(txt, openai_apikey):
    openai.api_key = openai_apikey
    path = str(time.time())
    import os
    os.mkdir(path)
    with open(path + '/doc.txt', mode = 'w', encoding = 'utf-8') as f:
        f.write(txt)
    
    documents = SimpleDirectoryReader(path).load_data()
    index = GPTVectorStoreIndex.from_documents(documents)
    template = (
        "你是一个有用的助手，能够根据图片内容准确地回答问题，并且尽量做到简洁明了:\n"
        "\n"
        "{context_str}\n"
        "\n"
        "{query_str}\n"
        "请你回复用户。\n"
    )
    qa_template = Prompt(template)
    query_engine = index.as_query_engine(text_qa_template = qa_template)
    return query_engine, gr.Textbox.update(visible = True), gr.Button.update(visible = True), gr.Markdown.update(
        value = '''建立索引成功，可以开始对话咯'''), gr.Chatbot.update(visible = True),0

    





def get_response(openai_apikey, msg, bot, query_engine):
    openai.api_key = openai_apikey
    
    query_str = '历史对话如下:\n'
    for his in bot:
        query_str += '用户：' + his[0] + '\n'
        query_str += 'bot:' + his[1] + '\n'
    query_str += '用户：' + msg + '\n'
    res = query_engine.query(query_str)
    print(res)
    bot.append([msg, str(res)])
    return bot[max(0, len(bot) - 3):]
        

def up_file(files):
    Doc_text_list = []
    for idx, file in enumerate(files):
        print(file.name)
        return gr.Textbox.update(visible = True), gr.Button.update(visible = True), gr.Markdown.update(value = '点击“建立索引”开始对话',)


with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            openai_apikey = gr.Textbox(label = 'openai api key', placeholder = '输入你的openai apikey')
            file = gr.File(file_types = ['.jpg'], label = '点击上传图片，格式需为jpg', file_count = 'multiple')
            txt = gr.Textbox(label = '图片解析结果', visible = False)
            with gr.Row():
                index_bu = gr.Button(value = '点击建立索引', visible = False)
            query_engine = gr.State([])
        with gr.Column():
            md = gr.Markdown('点击左侧区域上传图片进行解析')
            chat_bot = gr.Chatbot(visible = False)
            msg_txt = gr.Textbox(value = '消息框', placeholder = '输入信息发送', visible = False)
            chat_bu = gr.Button(visible = False)

    file.change(up_file, [file], [txt, index_bu, md])
    index_bu.click(doc_index, [openai_apikey, txt],[query_engine, msg_txt, chat_bu, md, chat_bot])
    chat_bu.click(get_response, [openai_apikey, msg_txt, chat_bot, query_engine], [chat_bot])


if __name__ == "__main__":
    demo.queue().launch()