import json
import requests
import os

# ====== 请在这里设置你的参数 ======
API_KEY = "sk-t2TTC9Wfx1U5j6EtnrON9mtxcbjCK7gb1vUd4qsSoEd5voEt"
MODEL_NAME = "deepseek-v3"
INPUT_PATH = "/home2/luotianyao/Lrm_Judge_192/0702data/Claude-3.7-sonnet-thinking/only_thinking/writing.json"
OUTPUT_PATH = "/home2/luotianyao/Lrm_Judge_192/0702data/Claude-3.7-sonnet-thinking/only_thinking/think-response-result/writing.json"
# ===================================

def query_deepseek(content, model_name, api_key, mode="thinking"):
    url = "https://4.0.wokaai.com/v1/chat/completions"
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json"
    }
    
    if mode == "thinking":
        instruction = "请根据以下thinking内容判断它更倾向于认为哪个回答更好，更好的回答的标准是：1.是两个回答都对的情况下，该回答的优点比另一个回答多，那么该回答更好；2.两个回答都错的情况下，该回答比另一个回答错误的地方更少，该回答更好；3.正确的回答永远比错误的回答更好；4.如果两个回答得一样好或一样差，那么回复A=B；5.其余情况严格回复A>B或B>A"
    else:
        instruction = "请根据以下thinking内容判断它更倾向于认为哪个回答更好，更好的回答的标准是：1.是两个回答都对的情况下，该回答的优点比另一个回答多，那么该回答更好；2.两个回答都错的情况下，该回答比另一个回答错误的地方更少，该回答更好；3.正确的回答永远比错误的回答更好；4.如果两个回答得一样好或一样差，那么回复A=B；5.其余情况严格回复A>B或B>A"
    
    prompt = f"""
{instruction}

内容:
{content}
    """.strip()

    data = {
        "model": model_name,
        "messages": [{"role": "user", "content": prompt}]
    }

    response = requests.post(url, headers=headers, json=data)
    response.raise_for_status()
    result = response.json()['choices'][0]['message']['content'].strip()

    print(f"🔹 DeepSeek 返回内容: {result}")

    if "A>B" in result:
        return "A>B"
    elif "B>A" in result:
        return "B>A"
    elif "A=B" in result:
        return "A=B"
    else:
        return "无法判断"

def main():
    # 加载数据
    with open(INPUT_PATH, 'r', encoding='utf-8') as f:
        data = json.load(f)

    if isinstance(data, dict):
        data = [data]
    elif not isinstance(data, list):
        raise ValueError("输入 JSON 必须是对象或对象列表")

    results = []

    for item in data:
        # 处理 judge_thinking
        judge_thinking = item.get("judge_thinking", "")
        if judge_thinking:
            try:
                result_thinking = query_deepseek(judge_thinking, MODEL_NAME, API_KEY, mode="thinking")
                item["judge_thinking_result"] = result_thinking
            except Exception as e:
                print(f"❌ 请求 judge_thinking 出错: {e}")
                item["judge_thinking_result"] = "请求失败"
        else:
            print("⚠️ 缺少 judge_thinking，跳过该字段")
            item["judge_thinking_result"] = "缺少数据"

        # 处理 judge_response
        judge_response = item.get("judge_response", "")
        if judge_response:
            try:
                result_response = query_deepseek(judge_response, MODEL_NAME, API_KEY, mode="response")
                item["judge_response_result"] = result_response
            except Exception as e:
                print(f"❌ 请求 judge_response 出错: {e}")
                item["judge_response_result"] = "请求失败"
        else:
            print("⚠️ 缺少 judge_response，跳过该字段")
            item["judge_response_result"] = "缺少数据"

        results.append(item)

    # 保存结果
    os.makedirs(os.path.dirname(OUTPUT_PATH), exist_ok=True)
    with open(OUTPUT_PATH, 'w', encoding='utf-8') as f:
        json.dump(results, f, ensure_ascii=False, indent=4)

    print(f"\n✅ 处理完成，结果已保存到: {OUTPUT_PATH}")

if __name__ == "__main__":
    main()
