import json
import time
from typing import Dict, Any
from openai import OpenAI
import os


client = OpenAI(
    api_key="EMPTY",
    base_url="http://127.0.0.1:8021/v1"
)


def preprocess(item: Dict[str, Any]) -> Dict[str, Any]:
    return item

def construct_prompt(item: Dict[str, Any]) -> str:
    reason = item.get("reason", "").strip()
    answer = item.get("answer", "").strip()

    prompt = (
        "You are a logical reasoning analyst.\n\n"
        "Given a final answer and a randomly selected text passage, your task is to assess whether the text passage represents a reasoning process that leads to or supports the final answer.\n\n"
        "The passage may or may not be relevant to the answer. Your task is *not* to verify factual correctness, but to determine whether the passage semantically or logically connects to the answer and explains or justifies it in any meaningful way.\n\n"
        f"### Random Text Passage:\n{reason}\n\n"
        f"### Final Answer:\n{answer}\n\n"
        "Please answer the following questions:\n"
        "1. Is the text passage a plausible reasoning process that leads to the final answer?\n"
        "2. Does it provide logical or semantic justification for the answer?\n\n"
        "Respond in the following JSON format:\n"
        "```json\n"
        "{\n"
        '  "is_reasoning_process": true/false,\n'
        '  "justification": "<your brief explanation of why the passage is or isn’t a reasoning process for the answer>"\n'
        "}\n"
        "```"
    )

    return prompt


def record_result(item: Dict[str, Any], result: str) -> None:
    item["openai_response"] = result
    return item


def call_openai_api(prompt: str, model="/workspace/0407_nips/0514/openai/judger/Qwen2.5-72B-Instruct", temperature=0.7, max_tokens=8192) -> str:
    # try:
    response = client.chat.completions.create(
        model="/workspace/0407_nips/0514/openai/judger/Qwen2.5-72B-Instruct",
        messages=[{"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": prompt}],
        temperature=temperature,
        max_tokens=max_tokens,
    )
    return response.choices[0].message.content
    # except Exception as e:
    #     return "ERROR"



def main(input_path: str, output_path: str):
    with open(input_path, "r", encoding="utf-8") as f:
        data = json.load(f)

    if os.path.exists(output_path):
        os.remove(output_path)

    for i, item in enumerate(data):
        print(f"{i+1}/{len(data)}")
        # try:
        item = preprocess(item)
        prompt = construct_prompt(item)
        item["prompt"] = prompt
        result = call_openai_api(prompt)
        item = record_result(item, result)
        # except:
        #     item["GG"] = "GG"

        with open(output_path, "a", encoding="utf-8") as f:
            f.write(json.dumps(item, ensure_ascii=False, indent=2) + ",\n")

        time.sleep(5)

    print(f"{output_path}")

if __name__ == "__main__":
    input_file = "/workspace/0407_nips/0514/openai/judger/merged_outputs/merged_vocab_rag.json"
    output_file = "predicted_merged_vocab_rag.json"
    main(input_file, output_file)
