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

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

OUTPUT_DIR = "outputs"
os.makedirs(OUTPUT_DIR, exist_ok=True)

# 拆解 JD 成任务，并生成三个解决方案（增强容错）
def generate_solutions_from_jd(jd):
    message = f"""你是一个岗位分析助手，请根据以下JD内容：\n
1. 首先识别一个可以用来测试岗位候选人核心能力的具体任务；
2. 然后为这个任务生成三个不同的解决方案建议。

请严格以如下格式输出：
任务:（任务描述）\n方案1:（内容）\n方案2:（内容）\n方案3:（内容）\n
JD: {jd}
"""
    response = taskgen_client.predict(message=message, api_name="/chat")

    task_match = re.search(r"任务[:：](.*)", response)
    solutions = re.findall(r"方案[123][:：](.*)", response)

    task = task_match.group(1).strip() if task_match else "任务解析失败"
    s1 = solutions[0].strip() if len(solutions) > 0 else "方案1解析失败"
    s2 = solutions[1].strip() if len(solutions) > 1 else "方案2解析失败"
    s3 = solutions[2].strip() if len(solutions) > 2 else "方案3解析失败"

    return task, s1, s2, s3

# 构建 Gradio UI
def build_ui():
    with gr.Blocks() as demo:
        gr.Markdown("## 📌 JD 任务拆解 + 解决方案选择 Demo")

        jd_input = gr.Textbox(label="输入 JD", placeholder="请输入岗位描述 JD")
        task_output = gr.Textbox(label="拆解出的测试任务", interactive=False)

        generate_button = gr.Button("🚀 拆解任务并生成方案")

        sol1 = gr.Textbox(label="方案1", interactive=False)
        sol2 = gr.Textbox(label="方案2", interactive=False)
        sol3 = gr.Textbox(label="方案3", interactive=False)

        select_radio = gr.Radio(choices=["1", "2", "3"], label="请选择你最满意的解决方案编号")
        user_solution = gr.Textbox(lines=4, label="📄 或者填写你自己的解决方案（可选）")
        comment = gr.Textbox(lines=3, label="📝 请选择的理由或批注")

        submit = gr.Button("✅ 提交 RLHF 数据")
        feedback = gr.Textbox(label="系统反馈", interactive=False)

        task_state = gr.State()

        def handle_jd_input(jd_text):
            task, s1, s2, s3 = generate_solutions_from_jd(jd_text)
            return task, s1, s2, s3, task

        def handle_submit(selected_idx, user_input_text, comment_text, task_text):
            record = {
                "task": task_text,
                "selected_index": selected_idx,
                "user_solution": user_input_text,
                "comment": comment_text,
                "timestamp": datetime.now().isoformat()
            }
            try:
                with open("rlhf_jd_data.jsonl", "a", encoding="utf-8") as f:
                    json.dump(record, f, ensure_ascii=False)
                    f.write("\n")
                return f"✅ 数据已保存，选择方案 {selected_idx}"
            except Exception as e:
                return f"❌ 保存失败：{str(e)}"

        generate_button.click(fn=handle_jd_input, inputs=[jd_input], outputs=[task_output, sol1, sol2, sol3, task_state])
        submit.click(fn=handle_submit, inputs=[select_radio, user_solution, comment, task_state], outputs=[feedback])

    return demo

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