import openai, tiktoken, os
import pandas as pd
from openai.embeddings_utils import get_embedding, cosine_similarity
import gradio as gr
import pickle
from utils import *

openai.api_key = os.getenv("OPENAI_API_KEY")

# Split the text into chunks of 200 Chinese characters
def split_chinese_text_into_chunks(text, chunk_size=200):
    chunks = [text[i:i+chunk_size] for i in range(0, len(text), chunk_size)]
    return chunks

# Create embeddings
def create_embedding(text1, text2, text3, text4, text5, text6, text7, text8, text9, text10):
    # Concatenated texts
    concatenated_text = text1 + text2 + text3 + text4 + text5 + text6 + text7 + text8 + text9 + text10
    # Get the chunks
    chunks = split_chinese_text_into_chunks(concatenated_text)
    # Convert chunks into a Pandas DataFrame
    df = pd.DataFrame({"text": chunks})
    # Removing any row with empty text
    df = df[df.text.ne('')]
    # Counting the number of tokens for each text 
    df["n_tokens"] = df.text.apply(lambda x: len(encoding.encode(str(x))))
    # filter too long text if any
    df = df[df.n_tokens <= max_tokens]
    df["embedding"] = df.text.apply(lambda x: get_embedding(x, engine=embedding_model))
    print('Start saving pkl file')
    with open('tmp.pkl', 'wb') as f:
        pickle.dump(df, f)
    print('Finish saving pkl file')       
    return "文章分析完成囉！請到頁面最上方，切換到 Chatbot 頁面進行問答或是文章生成"

def Bot(prompt):
    with open('tmp.pkl', 'rb') as f:
        df = pickle.load(f)

    prompt_embedding = get_embedding(prompt, engine=embedding_model)
    df["similarity"] = df.embedding.apply(lambda x: cosine_similarity(x, prompt_embedding))
    results = (df.sort_values("similarity", ascending=False))
    system = """
    你是一個萬能文字助手，你擅長從大量的文章中，辨識出相關主題，並整理成重點摘要。
    """
    messages = [{"role": "system", "content": system},]
    messages.append({"role": "user", "content": prepare_prompt(prompt, results)})
    return answer(messages)

with gr.Blocks() as demo:
    gr.Markdown(
    """
    # 文章彙整工具
    輸入希望快速吸收資訊來源，即可使用問答的方式獲取資訊！

    請輸入 Open AI API key 以使用本服務
    """
    )
    openai_api_key_textbox = gr.Textbox(placeholder="Paste your OpenAI API key (sk-...) and hit Enter",
                                    show_label=False, lines=1, type='password')
    openai_api_key_textbox.change(set_openai_api_key,
                                    inputs=[openai_api_key_textbox],
                                    outputs=[])
    openai_api_key_textbox.submit(set_openai_api_key,
                                    inputs=[openai_api_key_textbox],
                                    outputs=[])

    with gr.Tab("Articles"):
        gr.Markdown("請在下列文字匡中輸入待彙整的文章")
        input1 = gr.Textbox(label="Articles 1")
        input2 = gr.Textbox(label="Articles 2")
        input3 = gr.Textbox(label="Articles 3")
        input4 = gr.Textbox(label="Articles 4")
        input5 = gr.Textbox(label="Articles 5")
        input6 = gr.Textbox(label="Articles 6")
        input7 = gr.Textbox(label="Articles 7")
        input8 = gr.Textbox(label="Articles 8")
        input9 = gr.Textbox(label="Articles 9")
        input10 = gr.Textbox(label="Articles 10")
        output_text = gr.Textbox(label="處理進度")
        text_button = gr.Button("開始分析")
        text_button.click(fn=create_embedding, inputs=[input1, input2, input3, input4, input5, input6, input7, input8, input9, input10], outputs=output_text, api_name="create_embedding")

    with gr.Tab("Chatbot"):
        def predict(message, history):
            bot_message = Bot(prompt = message)

            partial_message = ""
            for chunk in bot_message:
                if len(chunk['choices'][0]['delta']) != 0:
                    partial_message = partial_message + chunk['choices'][0]['delta']['content']
                    yield partial_message
 
        gr.ChatInterface(predict)

demo.queue().launch()