# -*- coding: utf-8 -*-

import gradio as gr
from utils import prompt
from utils import engine
from utils.datasetSaver import login
import time
MODEL_NAME = 'gpt-3.5-turbo'
# MODEL_NAME = 'gpt-4-turbo-preview'
# MODEL_NAME = 'glm-4'

# 自定义system
# 设置端口号，默认7560，遇冲突可自定义
SERVER_PORT = 7562


story_path = r"/home/yangkelang/program/novalConvertProject/data/stories/十八岁出门远行.txt"

prompt_text = prompt.prompt_generator(story_path)


# 调用gpt的bot
LLM = engine.ChatGPT(model=MODEL_NAME,init_system={"role": "system", "content": prompt_text})
# LLM = engine.zhiPuGlm(model=MODEL_NAME, init_system={"role": "system", "content": prompt_text})

initial_response = LLM.get_response()

def predict(input, chatbot):
    """ 调用openai接口，获取答案
    """
    chatbot.append((input, ""))
    # 找chatgpt要答案
    response = LLM.get_response(input)
    characters = ''
    for character  in response:
        characters +=character
        chatbot[-1] = (input,characters)
        time.sleep(0.05)
        yield chatbot
    # return chatbot

def reset_user_input():
    return gr.update(value='')

def reset_user_input_new_game():
    return gr.update(value='newgame')

def reset_state():
    LLM.clean_history()
    return []

def save_history():
    LLM.save_history()
    return []

def new_game():
    return gr.update(value=[[None,initial_response]])


def main():
    with gr.Blocks() as demo:
        gr.HTML("""<h1 align="center">{}</h1>""".format(MODEL_NAME))
        # gradio的chatbot

        chatbot = gr.Chatbot(value=[[None,initial_response]])

        with gr.Row():
            with gr.Column(scale=4):
                with gr.Column(scale=50):
                    user_input = gr.Textbox(show_label=False, placeholder="Input...",container=False)
                with gr.Column(min_width=32, scale=1):
                    submitBtn = gr.Button("Submit", variant="primary")
            with gr.Column(scale=1):
                newBtn = gr.Button("新的游戏")
        # 提交问题
        submitBtn.click(predict, [user_input, chatbot],
                        [chatbot], show_progress=True)
        submitBtn.click(reset_user_input, [], [user_input])
        # 新的游戏

        newBtn.click(reset_user_input_new_game,[],[user_input])
        newBtn.click(reset_state, outputs=[chatbot], show_progress=True).then(save_history, outputs=[chatbot], show_progress=True).then(new_game,[],[chatbot])
        # newBtn.click(predict, [user_input, chatbot],[chatbot], show_progress=True)
        newBtn.click(reset_user_input, [], [user_input])


    # demo.queue().launch(share=True, inbrowser=True, server_port=SERVER_PORT,auth=login)
    demo.queue().launch(share=True, inbrowser=True, auth=login)
    
#chmod +x /home/yangkelang/lib/python3.10/site-packages/gradio/frpc_linux_amd64_v0.2


if __name__ == '__main__':
    main()