import json
import os
from openai import OpenAI
import argparse

parser = argparse.ArgumentParser(description="Example script")
parser.add_argument("--batch_id", type=str, required=True, help="Input file path")
args = parser.parse_args()

os.environ["OPENAI_API_KEY"] = "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"

client = OpenAI()

OUTPUT_DIR = f"batchapi/data_all/jobs_2/{args.batch_id}"

with open(f"batchapi/jobs/jobs_2/{args.batch_id}", "r") as f:
    # BATCH_ID = "batch_69899f6279a48190abaad4d65407768b"
    BATCH_ID = f.readline().strip()

os.makedirs(OUTPUT_DIR, exist_ok=True)

batch = client.batches.retrieve(BATCH_ID)
output_file_id = batch.output_file_id

result = client.files.content(output_file_id).text

for line in result.splitlines():
    record = json.loads(line)

    img_name = record["custom_id"]
    json_name = img_name.rsplit(".", 1)[0] + ".json"
    out_path = os.path.join(OUTPUT_DIR, json_name)

    if os.path.exists(out_path):
        continue

    # 🔥 Correct way to extract text output
    response_body = record["response"]["body"]
    output_text = response_body["output"][0]["content"][0]["text"]

    # Try to parse the model's JSON text
    try:
        parsed_json = json.loads(output_text)
    except json.JSONDecodeError:
        print(f"Invalid JSON for {img_name}")
        print(output_text)
        continue

    with open(out_path, "w") as f:
        json.dump(parsed_json, f, indent=2)

print("All results saved.")