import gradio as gr

import os
from ModelMerge.models import chatgpt

GPT_ENGINE = os.environ.get('GPT_ENGINE', 'gpt-4o')
API = os.environ.get('API', None)

systemprompt = (
    "你是一位旅行规划专家。你需要帮助用户规划旅行行程，给出合理的行程安排。"
    "- 如果用户提及要从一个城市前往另外一个城市，必须使用 get_Round_trip_flight_price 查询两个城市半年内往返机票价格信息。"
    "- 在规划行程之前，必须使用 get_city_tarvel_info 查询城市的景点旅行攻略信息。"
    "- 查询攻略后，你需要分析用户个性化需求。充分考虑用户的年龄，情侣，家庭，朋友，儿童，独自旅行等情况。排除不适合用户个性化需求的景点。之后输出符合用户需求的景点。"
    "- 综合用户游玩时间，适合用户个性化需求的旅游城市景点，机票信息和预算，给出真实准确的旅游行程，包括游玩时长、景点之间的交通方式和移动距离，每天都要给出总的游玩时间。"
    "- 根据查到的景点介绍结合你自己的知识，每个景点必须包含你推荐的理由和景点介绍。介绍景点用户游玩的景点，景点介绍尽量丰富精彩，吸引用户眼球，不要直接复述查到的景点介绍。"
    "- 每个景点都要标注游玩时间、景点之间的交通方式和移动距离还有生动的景点介绍。必须标注景点之间的交通方式和移动距离还有景点游玩时长，否则不要安排该景点。"
    "- 尽量排满用户的行程，不要有太多空闲时间。"
    "- 景点的描述必须遵循以下格式："
    "{景点名称}（{游玩时间}：{景点介绍}。"
    "{下一个地点的交通方式}，{下一个地点的移动距离}"
    "{下一个景点名称}（{游玩时间}：{景点介绍}。"
)

chatgptbot = chatgpt(api_key=API, engine=GPT_ENGINE, system_prompt=systemprompt)

with gr.Blocks(fill_height=True) as demo:
    with gr.Column():
        chatbot = gr.Chatbot(show_label=False, elem_id="chatbox", scale=10, height=900)  # 设置聊天框高度
        with gr.Row():
            msg = gr.Textbox(placeholder="输入你的问题...", elem_id="inputbox", scale=10)
            clear = gr.Button("清除", elem_id="clearbutton")

    # 用户输入处理函数，记录用户的问题
    def user(user_message, history):
        return "", history + [[user_message, None]]

    # 机器人回答函数，根据用户的问题生成回答
    def bot(history):
        print(history)
        user_message = history[-1][0]
        history[-1][1] = ""
        answer = ""
        for text in chatgptbot.ask_stream(user_message):
            print(text, end="")
            if "🌐" in text:
                history[-1][1] = text
            else:
                answer += text
                history[-1][1] = answer
            yield history

    # 提交用户消息和处理回答
    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
        bot, chatbot, chatbot
    )
    # 清除聊天记录
    clear.click(lambda: None, None, chatbot, queue=False)

demo.launch()