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

# 初始化任务生成客户端（腾讯混元 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)
        testable_skills = 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 提取任务")
        generate_solution = gr.Button("🎨 生成解法选项")
        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):
            # 步骤 1：解析任务，并生成具体的可测试任务
            message = f"""你是一个任务解析助手，请根据以下 JD 提取出一个可以测试硬技能的任务，并具体列出需要的硬技能：
            JD: {jd_text}
            请提取出以下内容：
            1. 任务描述：可以通过具体操作测试的任务
            2. 硬技能：为测试该任务所需的硬技能（例如编程语言、工具、技术等）"""
            
            response = taskgen_client.predict(message=message, api_name="/chat")
            
            # 任务描述和可测试硬技能的提取
            task_match = re.search(r"任务描述[:：](.*?)(、|。)", response)
            skills_match = re.search(r"硬技能[:：](.*?)(、|。)", response)
            
            task = task_match.group(1).strip() if task_match else "任务解析失败"
            testable_skills = skills_match.group(1).strip() if skills_match else "无硬技能信息"
            
            return task, testable_skills, jd_text, task, testable_skills

        def handle_generate_solution(task):
            # 步骤 2：生成解法选项
            solution_message = f"""根据以下任务描述，生成三个解决方案选项：
            任务描述：{task}
            请为该任务生成三个具体的解决方案，并列出其优缺点。"""
            
            solution_response = taskgen_client.predict(message=solution_message, api_name="/chat")
            
            # 解析解法选项（假设模型返回每个选项单独一行）
            solution_options = solution_response.split("\n")  # 假设解法选项按行分开
            return 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, testable_skills, jd_state, task_state])
        generate_solution.click(fn=handle_generate_solution, inputs=[task], outputs=[solution_options])
        submit.click(fn=handle_submit, inputs=[solution_options, comment, final_solution], outputs=[feedback])

    return demo

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