from agent import agent
import json


QUESTION_NUMBER = 18

with open("questions.json", "r") as f:
    question = json.load(f)[QUESTION_NUMBER]

response = agent.invoke(
    {
        "messages": f"{question['question']}\nFile name: {question['file_name']}" if question["file_name"] else question['question'],
    }
)['messages'][-1].content

response = response.split("</think>")[1].strip().split("FINAL ANSWER:")[-1].strip()
print(f"\n{response}")

with open("answers_payload.json", "r") as f:
    answers_payload = json.load(f)

answers_payload[QUESTION_NUMBER] = {
    "task_id": question["task_id"],
    "submitted_answer": response
}

with open("answers_payload.json", "w") as f:
    json.dump(answers_payload, f, indent=4)

with open("results_log.json", "r") as f:
    results_log = json.load(f)

results_log[QUESTION_NUMBER] = {
    "Task ID": question["task_id"],
    "Question": question["question"],
    "Submitted Answer": response
}

with open("results_log.json", "w") as f:
    json.dump(results_log, f, indent=4)