import gradio as gr
import openai
import time

# **配置 OpenAI API Key**
OPENAI_API_KEY = "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"  # 👈 **替换这里！**

# **使用新版本的 OpenAI Client**
client = openai.OpenAI(api_key=OPENAI_API_KEY)

def generate_task_structure(prompt):
    """
    通过 OpenAI GPT-4o 生成任务管理计划，**支持进度条 & 流式输出**
    """
    
    structured_prompt = f"""
    你是 `StratAI`，一个智能任务管理 AI 助手。请根据用户的需求，生成一份完整的任务管理计划。
    
    任务应包括：
    - **目标**
    - **主要任务**
    - **子任务**
    - **里程碑**
    - **时间安排**
    
    用户的需求是：{prompt}

    请输出结构化的任务规划，使其适合直接阅读。
    """

    # **1️⃣ 任务开始，初始化进度条**
    progress_steps = [
        ("🟢 开始处理任务...", 0.1),
        ("🟡 调用 OpenAI GPT-4o...", 0.3),
        ("🔵 AI 生成任务内容中...", 0.6),
        ("✅ 任务完成！", 1.0)
    ]

    # **2️⃣ 调用 OpenAI API（流式输出）**
    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {"role": "system", "content": "你是一个任务管理 AI 助手，帮助用户创建任务计划。"},
            {"role": "user", "content": structured_prompt}
        ],
        temperature=0.5,
        stream=True  # **开启流式响应**
    )

    # **3️⃣ UI 进度条**
    with gr.update() as ui:
        for text, progress in progress_steps:
            ui.progress(progress)  # **更新进度条**
            ui.update(text)  # **更新进度信息**
            time.sleep(0.5)

    # **4️⃣ 流式输出**
    full_text = ""
    for chunk in response:
        text_chunk = chunk.choices[0].delta.content or ""
        full_text += text_chunk
        yield full_text  # **实时输出**
    
    yield "✅ 任务已完成！"

# **Gradio UI**
with gr.Blocks() as demo:
    gr.Markdown("# 🧠 StratAI - AI 任务管理助手")
    gr.Markdown("输入你的想法，让 AI 生成 **目标 -> 任务 -> 里程碑 -> 排期**")

    user_input = gr.Textbox(label="输入你的想法", placeholder="例如：我想开一家精品咖啡店")

    with gr.Row():
        progress_bar = gr.Slider(minimum=0, maximum=1, step=0.01, label="任务进度", value=0, interactive=False)
    
    output_text = gr.Textbox(label="任务规划", lines=15)

    submit_button = gr.Button("生成任务")
    submit_button.click(fn=generate_task_structure, inputs=user_input, outputs=[progress_bar, output_text])

demo.launch()