import gradio as gr
from utils.openai_api import get_reply, get_tokens_length
from utils.read_file import read_xlsx_file, process_data

sessions = {}


def get_session(key):
    if key not in sessions:
        sessions[key] = {"count_now": 0, "reply": ""}
    return sessions[key]


def submit(action, data_source, file, max_length, message_first, message_after,
           session_key):
    session = get_session(session_key)
    count_now = session["count_now"]
    reply = session["reply"]
    if action == "预览消息":
        data = read_xlsx_file(data_source, file.name)
        if count_now == 0:
            message, count = process_data(data_source, max_length, data,
                                          message_first, len(data))
        else:
            message, count = process_data(data_source,
                                          max_length,
                                          data,
                                          message_after,
                                          len(data),
                                          count_now,
                                          reply=reply)
        return get_tokens_length(message), message, ""
    elif action == "获取最终回复(耗时较长)":
        while True:
            data = read_xlsx_file(data_source, file.name)
            if count_now == 0:
                message, count = process_data(data_source, max_length, data,
                                              message_first, len(data))
            else:
                message, count = process_data(data_source,
                                              max_length,
                                              data,
                                              message_after,
                                              len(data),
                                              count_now,
                                              reply=reply)
            count_now += count
            if count == 0:
                break
            else:
                print(message)
                reply = get_reply(message)
                print(reply)
        session["count_now"] = count_now
        session["reply"] = reply
        return get_tokens_length(message), "", reply
    elif action == "获取一轮回复":
        data = read_xlsx_file(data_source, file.name)
        if count_now == 0:
            message, count = process_data(data_source, max_length, data,
                                          message_first, len(data))
        else:
            message, count = process_data(data_source,
                                          max_length,
                                          data,
                                          message_after,
                                          len(data),
                                          count_now,
                                          reply=reply)
        count_now += count
        reply = get_reply(message)
        session["count_now"] = count_now
        session["reply"] = reply
        return get_tokens_length(message), message, reply


iface = gr.Interface(
    fn=submit,
    inputs=[
        gr.inputs.Dropdown(choices=["预览消息", "获取最终回复(耗时较长)", "获取一轮回复"],
                           label="操作",
                           default="预览消息"),
        gr.inputs.Dropdown(choices=["七麦", "点点"], label="数据来源", default="七麦"),
        gr.inputs.File(label="上传 xlsx 文件"),
        gr.inputs.Number(default=2048, label="长度限制(中文建议最大4096)"),
        gr.inputs.Textbox(
            lines=2,
            label="第一轮预定义消息",
            default=
            "接下来输入第{comment_num_start}条-第{comment_num_end}条app store中对一款app的评论(共{all_num}条)，请分条总结这款app的好评论与坏评论(各十条)(并在每条后面按照百分比给出这个观点的在好/坏评论中的当前占比),{data_string}"
        ),
        gr.inputs.Textbox(
            lines=2,
            label="后续预定义消息",
            default=
            "{reply_before},以上是对前{comment_num}条的分析结果，接下来分段输入第{comment_num_start}条-第{comment_num_end}条app store中对一款app的评论(共{all_num}条)，请分条总结这款app的好评论与坏评论(各十条)(并在每条后面按照百分比给出这个观点的在好/坏评论中的当前占比),{data_string}"
        ),
        gr.inputs.Textbox(lines=1,
                          label="session_key",
                          default="new_session_id(用于记录会话))"),
    ],
    outputs=[
        gr.outputs.Textbox(label="Token长度"),
        gr.outputs.Textbox(label="GPT输入的字符串（Message）"),
        gr.outputs.Textbox(label="GPT输出的字符串（Reply）"),
    ],
    layout="vertical",
    title="GPT4Comment",
    description="上传xlsx文件并输入预定义的消息，然后点击发送使用GPT",
)

iface.launch(debug=True)
# iface.launch(server_name="0.0.0.0", server_port=7881)