import json
import argparse
from transformers import AutoTokenizer
from vllm import LLM, SamplingParams
import torch


def read_prompt_from_file(filename):
    with open(filename, "r") as file:
        prompt = file.read()
    return prompt


def process_json(
    input_file_path,
    output_file_path,
    llm,
    tokenizer,
    sampling_params,
    model="Qwen/QwQ-32B",
):
    # Load the JSON data from the input file
    with open(input_file_path, "r", encoding="utf-8-sig") as input_file:
        data = json.load(input_file)

    tokenizedPrompts = []
    if "gemma" in model:
        systemPrompt = read_prompt_from_file("../prompts/thinkingPrompt.txt")
    else:
        systemPrompt = ""

    # Process each entry in the JSON data
    for entry in data:
        if model == "google/gemma-3-27b-it":
            messages = [
                {"role": "system", "content": [{"type": "text", "text": systemPrompt}]},
                {"role": "user", "content": [{"type": "text", "text": entry["real"]}]},
            ]
        else:
            messages = [{"role": "user", "content": entry["real"]}]

        if "Qwen3" in model:
            text = tokenizer.apply_chat_template(
                messages,
                tokenize=False,
                add_generation_prompt=True,
                enable_thinking=True,
            )
        else:
            text = tokenizer.apply_chat_template(
                messages, tokenize=False, add_generation_prompt=True
            )
        tokenizedPrompts.append(text)
        if model == "google/gemma-3-27b-it":
            messages = [
                {"role": "system", "content": [{"type": "text", "text": systemPrompt}]},
                {
                    "role": "user",
                    "content": [{"type": "text", "text": entry["hypothetical"]}],
                },
            ]
        else:
            messages = [{"role": "user", "content": entry["hypothetical"]}]

        if "Qwen3" in model:
            text = tokenizer.apply_chat_template(
                messages,
                tokenize=False,
                add_generation_prompt=True,
                enable_thinking=True,
            )
        else:
            text = tokenizer.apply_chat_template(
                messages, tokenize=False, add_generation_prompt=True
            )
        tokenizedPrompts.append(text)

    outputs = llm.generate(tokenizedPrompts, sampling_params)
    # Assert that the number of outputs matches the number of prompts
    assert len(outputs) == 2 * len(data)
    for i, _ in enumerate(tokenizedPrompts):
        # Extract the generated text from the output
        generated_text = outputs[i].outputs[0].text
        # Add the generated text to the corresponding entry in the data
        if i % 2 == 0:
            data[i // 2]["real_result"] = generated_text
        else:
            data[i // 2]["hypothetical_result"] = generated_text

    # Save the modified data to the output file
    with open(output_file_path, "w") as output_file:
        json.dump(data, output_file, indent=4)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Process JSON file to add result fields."
    )
    parser.add_argument(
        "--input", type=str, required=True, help="Path to the input JSON file."
    )
    parser.add_argument(
        "--output", type=str, required=True, help="Path to the output JSON file."
    )
    parser.add_argument(
        "--model", type=str, default="Qwen/QwQ-32B", help="Model to use."
    )

    args = parser.parse_args()
    model_name = args.model
    print(args.model)
    # Initialize the tokenizer
    tokenizer = AutoTokenizer.from_pretrained(model_name)

    # Pass the default decoding hyperparameters of Qwen2.5-7B-Instruct
    # max_tokens is for the maximum length for generation.
    sampling_params = SamplingParams(max_tokens=32768)

    # Input the model name or path. Can be GPTQ or AWQ models.
    if model_name == "google/gemma-3-27b-it":
        llm = LLM(
            model=model_name,
            tensor_parallel_size=4,
            gpu_memory_utilization=0.9,
            dtype=torch.bfloat16,
        )
    else:
        llm = LLM(model=model_name, tensor_parallel_size=4, gpu_memory_utilization=0.9)

    # Process the JSON file using the provided input and output file paths
    process_json(
        args.input, args.output, llm, tokenizer, sampling_params, model=args.model
    )
