# -*- coding: utf-8 -*-
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.text_splitter import CharacterTextSplitter
from langchain.vectorstores import Chroma

import openai
import time
import gradio as gr
import os

# 输入 API KEY
os.environ["OPENAI_API_KEY"] = "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"



#读取PDF文件
def doc_read_pdf(file):
    # 读取PDF
    reader = PdfReader(file)
    # reader = PdfReader('.\data\資治通鑑全集_部分1.pdf')
    raw_text = ''
    for i, page in enumerate(reader.pages):
        text = page.extract_text()
        if text:
            raw_text += text
    return raw_text

# 读取txt文件
def doc_read_txt(file):
    with open(file, encoding='utf-8') as f:
        text = f.read()
    return text

#补全

#从开始到调用openai模型前的一些步骤，主要是文件读取和拆解
def doc_split(file):
    #分解文本
    text_splitter = CharacterTextSplitter(
        separator = "\n",
        chunk_size = 1200,
        chunk_overlap  = 100,
        length_function = len,
    )
    # texts = text_splitter.split_text(raw_text)
    texts = text_splitter.split_text(doc_read_txt(file))

    return texts

#将文本向量化
def doc_vectorize(texts):
    embeddings = OpenAIEmbeddings()
    docsearch = Chroma.from_texts(texts, embeddings, metadatas=[{"source": str(i)} for i in range(len(texts))])
    return docsearch

#文本被拆解储存在数组texts中
# raw_file = 
texts = doc_split("./资治通鉴_1_残缺.txt")
docsearch = doc_vectorize((texts))

def openai_reply(word1, word2, word3, temp, file):
    if file:
        texts = doc_split(file)
        docsearch = doc_vectorize((texts))
    words = word1 + "*****" + word2 + "*****" + word3
    # 文本相似查找，最终结果是一个列表
    docs = docsearch.similarity_search(words)
    reference_1 = docs[0].page_content
    reference_2 = docs[1].page_content
    reference = reference_1 + reference_2
    print(words)
    request = "请使用文言文帮我补全[" + words + "]"
    response = openai.ChatCompletion.create(
        model="gpt-3.5-turbo",
        messages=[
            {"role": "system",
            "content": f"""你是一个古汉语与中国历史专家，擅长补全古代文献。在后续的会话中，你需要补全我给你的残缺文本，
            这些残缺文本用[]包括，并且其中的几段残缺汉字用“*”来代替着，且数量不明。
            请你阅读并且理解下文背景资料，并且补全我待会在会话中给出的残缺文本。如果你在背景资料中找不到相关文字，请根据你对于背景资料和我给出的残缺文本的理解，
            使用《资治通鉴》的文言文风格自行补全残缺文本。请注意，不要改动我提供的残缺文本中非“*”的原文，并且一定要用文言文替代“*”！！！
            \n背景资料：\n{reference}
            """},
            {"role": "user", "content": request},
        ],
        max_tokens=512,
        n=1,
        stop=None,
        temperature=temp
    )
    print(response.choices[0].message['content'])
    shijian = time.strftime("%Y年%m月%d日%H点%M分",time.localtime())
    answer = response.choices[0].message['content']
    return answer, reference, shijian


# 以下是界面搭建
headline = '碎片化文本复原'
description = """请给出至多三条碎片文本，系统将会根据文献数据库尽可能进行理解和匹配，给出猜想。
当然，你也可以上传自己的TXT文件作为数据来源之一。结果仅供参考和启发。"""

with gr.Blocks() as demo:
    gr.Markdown(f'<center><h1>{headline}</h1></center>')
    gr.Markdown(description)

    with gr.Row():
        with gr.Group():
            raw_text_1 = gr.Textbox(label='在此输入碎片文本')
            raw_text_2 = gr.Textbox(label='在此输入碎片文本')
            raw_text_3 = gr.Textbox(label='在此输入碎片文本')
            temp = gr.Slider(minimum=0.0, maximum=2.0, value=0.3, label="无序程度（Temperature）")
            Title = gr.Textbox(label='在此输入提交的材料的标题（无需加《》）')
            file = gr.File(
                label='上传你的本地TXT文件', file_types=['.txt']
            )
            
            btn = gr.Button(value='提交')
            btn.style(full_width=True)

        with gr.Group():
            shijian = gr.Label(label='生成时时间')
            answer = gr.Textbox(label='回答')
            reference = gr.Textbox(label='参考材料')

        
        btn.click(
            openai_reply,
            inputs= [raw_text_1, raw_text_2, raw_text_3, temp, file],
            outputs= [answer, reference, shijian],
        )


demo.launch()

# if __name__ == "__main__":
#     demo.launch(server_port=7860, share=True)