import json
import re
import random
import math
import os
import time
from openai import OpenAI

client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

with open('datasets/gsm8k_test_with_code_1000.jsonl', 'r') as infile:
    data = []
    for line in infile:
        try:
            data.append(json.loads(line))
        except Exception as e:
            print(f"Error: {e}")

for entry_id, entry in enumerate(data):

    prompt = f"""
Your are given a math problem and a code that solves it. The function has several parameters.
Write a function call with the one parameter replaced by variable previous_answer, use the same values for other parameters as in the code provided.
Then replace the corresponding number in the question with the tag <variable>. Do not change anything in the question except replacing one value with <variable>.
Write the new question and the function call in the following format:
<function_call>
function_call
</function_call>

<modified_question>
modified_question
</modified_question>

Problem:

{entry['question']}

Code:

{entry['code_clean']}
    """

    print(f"Prompt: {prompt}")

    while True:
        try:
            completion = client.chat.completions.create(
                model="o3",
                messages=[
                    {"role": "developer", "content": "You are a helpful assistant."},
                    {"role": "user", "content": prompt}
                ],
            )
            response = completion.choices[0].message.content.strip()
            print(f"Response: {response}")
            function_call = response.split('<function_call>')[1].split('</function_call>')[0].strip()
            modified_question = response.split('<modified_question>')[1].split('</modified_question>')[0].strip()
            print(f"Function call: {function_call}")
            print(f"Modified question: {modified_question}")
            break
        except Exception as e:
            response = f"Error: {e}"
            function_call = None
            modified_question = None
            print(f"Error: {e}")

    # entry['response'] = response
    entry['function_call'] = function_call
    entry['modified_question'] = modified_question

    print(f"Processed {entry_id+1}/{len(data)} entries.")
    time.sleep(1)

    while True:
        try:
            with open('datasets/gsm8k_test_with_code_with_variables_in_questions_1000.jsonl', 'a') as outfile:
                json.dump(entry, outfile, ensure_ascii=False)
                outfile.write("\n")
            break
        except Exception as e:
            print(f"Error: {e}")
            time.sleep(10)
