import gradio as gr
from gradio_client import Client
import re
import json
from datetime import datetime

# 初始化任务生成客户端（腾讯混元 Space）
taskgen_client = Client("tencent/Hunyuan-Large")

# 构建 Gradio UI
def build_ui():
    with gr.Blocks() as demo:
        gr.Markdown("## 📝 JD 任务解析与解法选择 Demo")

        jd_text = gr.Textbox(label="职位描述 (JD)", placeholder="请输入职位描述文本，如‘负责销售数据分析与报告撰写’")
        task = gr.Textbox(label="提取的任务", interactive=False)
        solution_options = gr.CheckboxGroup(label="任务解法选项", choices=["选项1", "选项2", "选项3"], type="index")
        final_solution = gr.Textbox(label="您的最终解法", lines=3, placeholder="请提供您的解法")
        comment = gr.Textbox(label="📝 选择理由", lines=3, placeholder="为什么选择这个解法")

        generate_task = gr.Button("🔍 解析 JD 提取任务")
        submit = gr.Button("✅ 提交选择与解法（保存为RLHF数据）")
        feedback = gr.Textbox(label="系统反馈", interactive=False)

        jd_state = gr.State("")
        task_state = gr.State("")
        solution_state = gr.State([])

        def handle_jd_submit(jd_text):
            # 从 JD 中提取任务描述
            message = f"""你是一个任务解析助手，请根据以下 JD 提取出相关的任务：
            JD: {jd_text}
            请提取一个任务描述并用简洁明了的语言表达。
            """
            response = taskgen_client.predict(message=message, api_name="/chat")
            task_match = re.search(r"任务[:：](.*)", response)

            task = task_match.group(1).strip() if task_match else "任务解析失败"
            # 生成解法选项
            solution_options = [
                f"解法1：基于数据分析工具，生成详细报告",
                f"解法2：通过市场调研，提出优化建议",
                f"解法3：结合用户需求，制定定制化销售策略"
            ]

            return task, solution_options, jd_text, task, solution_options

        def handle_submit(selected_option, comment_text, final_solution_text):
            record = {
                "task": task_state,
                "solution_options": solution_state,
                "selected_option": selected_option,
                "comment": comment_text,
                "final_solution": final_solution_text,
                "timestamp": datetime.now().isoformat()
            }
            try:
                with open("rlhf_data.jsonl", "a", encoding="utf-8") as f:
                    json.dump(record, f, ensure_ascii=False)
                    f.write("\n")
                feedback_text = f"✅ RLHF 数据已保存，选择了解法 {selected_option}"
            except Exception as e:
                feedback_text = f"❌ 保存失败：{str(e)}"
            return feedback_text

        generate_task.click(fn=handle_jd_submit, inputs=[jd_text], outputs=[task, solution_options, jd_state, task_state, solution_state])
        submit.click(fn=handle_submit, inputs=[solution_options, comment, final_solution], outputs=[feedback])

    return demo

if __name__ == "__main__":
    demo = build_ui()
    demo.launch()