# Please install OpenAI SDK first: `pip3 install openai`

from openai import OpenAI
from tqdm import tqdm
import os
import json
import random

random.seed(42)
client = OpenAI(api_key="[API-KEY]", base_url="https://api.deepseek.com")
client_beta = OpenAI(api_key="[API-KEY]", base_url="https://api.deepseek.com/beta")

def dpsk(text):
    messages = [{"role": "user", "content": text}]
    response = client.chat.completions.create(
        model="deepseek-reasoner",
        messages=messages
    )

    reasoning_content = response.choices[0].message.reasoning_content
    content = response.choices[0].message.content
    return reasoning_content, content


def dpsk_prefix(text, prefix):
    messages = [
        {"role": "user", "content": text},
        {"role": "assistant", "content": prefix, "prefix": True}
    ]
    response = client_beta.chat.completions.create(
        model="deepseek-reasoner",
        messages=messages,
        max_tokens=10,
    )
    reasoning_content = response.choices[0].message.reasoning_content
    content = response.choices[0].message.content
    return reasoning_content, content

# question = "你是DeepSeek-R1还是V3？"
# reasoning, content = dpsk_prefix(question, "我是")
# print("REASONING")
# print(reasoning)
# print("CONTENT")
# print(content)
# exit(0)

# question = "你好！"
# reasoning, content = dpsk(question)
# print("REASONING")
# print(reasoning)
# print("CONTENT")
# print(content)
# exit(0)

files = [
    # "nature-cn.json",
    # "social-cn.json",
    "space-cn.json",
    # "time-cn.json",
    # "space+nature-cn.json",
]

direct_prompts = [
    "让我直接回答用户，不要思考过程。答案是：【",
    "我现在直接进行回答。我不应该输出思考过程。答案是：【",
    "我将会直接回答问题，不需要思考。答案是：【",
    "我不应该思考，我应该直接回答该问题。答案是：【",
]

input_directory = "path/to/LongCoT/Knowlogic/finaldata"
output_dir = "path/to/LongCoT/Knowlogic/dpsk_results"
os.makedirs(output_dir, exist_ok=True)

for file in files:
    print(file)
    try:
        with open(os.path.join(output_dir, file), "r", encoding='utf-8') as f:
            data = json.load(f)
    except:
        print("NO EXISTING RESULTS!")
        input_file = os.path.join(input_directory, file)    
        with open(input_file, 'r', encoding='utf-8') as f:
            data = json.load(f)

    for item in tqdm(data):
        if item.get("reasoning", None) != None:
            continue

        question = item['question']
        reasoning, answer = dpsk(question)

        item['reasoning'] = reasoning
        item['model_answer'] = answer
        direct_answers = []
        for prompt in direct_prompts:
            for _ in range(2):
                _, content = dpsk_prefix(question, prompt)
                direct_answers.append(content)
        item["direct_answers"] = direct_answers
        
        output_path = os.path.join(output_dir, file)
        with open(output_path, 'w', encoding='utf-8') as outfile:
            json.dump(data, outfile, ensure_ascii=False, indent=4)
