# gradio_demo.py

import gradio as gr
from orchestrator import Orchestrator, DialogueState

# --- 全局配置 ---
MODEL_NAME = "qwen2.5-72b-instruct" 

def create_orchestrator():
    """创建一个新的Orchestrator实例，方便重置。"""
    return Orchestrator(model=MODEL_NAME)

def process_message(user_input, chat_history, orchestrator_state):
    """处理用户输入的核心函数"""
    if orchestrator_state is None:
        orchestrator_state = create_orchestrator()

    response = orchestrator_state.process_user_message(user_input)
    chat_history.append({"role": "user", "content": user_input})
    chat_history.append({"role": "assistant", "content": response})
    print(chat_history)
    new_state_text = f"当前状态: {orchestrator_state.state.name}"
    
    return chat_history, orchestrator_state, new_state_text, ""

def reset_conversation():
    """
    【核心功能】重置会话的函数。
    它会创建一个全新的Orchestrator实例，并清空前端UI。
    """
    new_orchestrator = create_orchestrator()
    initial_state_text = f"当前状态: {new_orchestrator.state.name}"
    
    # 返回一个元组，用于清空和重置UI组件
    return [], new_orchestrator, initial_state_text, ""

# --- 构建Gradio界面 ---
with gr.Blocks(theme=gr.themes.Soft(), title="运筹优化对话平台") as demo:
    # gr.State 用于在会话中存储后端Orchestrator对象
    orchestrator_state = gr.State()

    gr.Markdown(
        """
        # 运筹优化智能对话平台
        通过对话收集信息，并在信息充足时自动切换到建模阶段。
        """
    )
    
    chatbot = gr.Chatbot(label="对话窗口", height=500, type="messages")
    
    state_display = gr.Textbox(
        label="系统状态", 
        value=f"当前状态: {DialogueState.REQUIREMENT_ELICITATION.name}",
        interactive=False
    )
    
    with gr.Row():
        user_input_textbox = gr.Textbox(
            label="输入你的需求",
            placeholder="例如：我想解决一个生产计划问题...",
            scale=4
        )
        send_btn = gr.Button("发送", variant="primary", scale=1)

    # 【这里是关键】 在界面上添加一个重置按钮
    reset_btn = gr.Button("🔄 新的对话 (重置)", variant="stop")

    # --- 绑定事件 ---
    
    # 准备好输入和输出组件列表，方便复用
    # 注意：这里的顺序必须和函数返回值的顺序一一对应
    outputs_list = [chatbot, orchestrator_state, state_display, user_input_textbox]
    
    # 绑定“发送”按钮和回车键
    send_btn.click(
        fn=process_message,
        inputs=[user_input_textbox, chatbot, orchestrator_state],
        outputs=outputs_list
    )
    user_input_textbox.submit(
        fn=process_message,
        inputs=[user_input_textbox, chatbot, orchestrator_state],
        outputs=outputs_list
    )
    
    # 【这里是关键】 绑定“重置”按钮的点击事件
    reset_btn.click(
        fn=reset_conversation,
        inputs=None, # 重置不需要从UI获取输入
        outputs=outputs_list # 重置后需要更新的组件列表和发送后一样
    )

# 【修改后】的代码
if __name__ == "__main__":
    demo.launch()